Official implementation of Generative Diffusion Models for Antibody Design, Docking, and Optimization.
import os
import shutil
import argparse
from diffab.tools.dock.hdock import HDockAntibody
from diffab.tools.runner.design_for_pdb import args_factory, design_for_pdb


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--antigen', type=str, required=True)
    parser.add_argument('--antibody', type=str, default='./data/examples/3QHF_Fv.pdb')
    parser.add_argument('--heavy', type=str, default='H', help='Chain id of the heavy chain.')
    parser.add_argument('--light', type=str, default='L', help='Chain id of the light chain.')
    parser.add_argument('--hdock_bin', type=str, default='./bin/hdock')
    parser.add_argument('--createpl_bin', type=str, default='./bin/createpl')
    parser.add_argument('-n', '--num_docks', type=int, default=10)
    parser.add_argument('-c', '--config', type=str, default='./configs/test/codesign_single.yml')
    parser.add_argument('-o', '--out_root', type=str, default='./results')
    parser.add_argument('-t', '--tag', type=str, default='')
    parser.add_argument('-s', '--seed', type=int, default=None)
    parser.add_argument('-d', '--device', type=str, default='cuda')
    parser.add_argument('-b', '--batch_size', type=int, default=16)
    args = parser.parse_args()

    hdock_missing = []
    if not os.path.exists(args.hdock_bin):
        hdock_missing.append(args.hdock_bin)
    if not os.path.exists(args.createpl_bin):
        hdock_missing.append(args.createpl_bin)
    if len(hdock_missing) > 0:
        print("[WARNING] The following HDOCK applications are missing:")
        for f in hdock_missing:
            print(f" > {f}")
        print("Please download HDOCK from http://huanglab.phys.hust.edu.cn/software/hdocklite/ "
                "and put `hdock` and `createpl` to the above path.")
        exit()

    antigen_name = os.path.basename(os.path.splitext(args.antigen)[0])
    docked_pdb_dir = os.path.join(os.path.splitext(args.antigen)[0] + '_dock')
    os.makedirs(docked_pdb_dir, exist_ok=True)
    docked_pdb_paths = []
    for fname in os.listdir(docked_pdb_dir):
        if fname.endswith('.pdb'):
            docked_pdb_paths.append(os.path.join(docked_pdb_dir, fname))
    if len(docked_pdb_paths) < args.num_docks:
        with HDockAntibody() as dock_session:
            dock_session.set_antigen(args.antigen)
            dock_session.set_antibody(args.antibody)
            docked_tmp_paths = dock_session.dock()
            for i, tmp_path in enumerate(docked_tmp_paths[:args.num_docks]):
                dest_path = os.path.join(docked_pdb_dir, f"{antigen_name}_Ab_{i:04d}.pdb")
                shutil.copyfile(tmp_path, dest_path)
                print(f'[INFO] Copy {tmp_path} -> {dest_path}')
                docked_pdb_paths.append(dest_path)

    for pdb_path in docked_pdb_paths:
        current_args = vars(args)
        current_args['tag'] += antigen_name
        design_args = args_factory(
            pdb_path = pdb_path,
            **current_args,
        )
        design_for_pdb(design_args)


if __name__ == '__main__':
    main()

AbDesign/design_dock.py
pengzhangzhi/ab_opt
from diffab.tools.runner.design_for_pdb import args_from_cmdline, design_for_pdb

if __name__ == '__main__':
    design_for_pdb(args_from_cmdline())

AbDesign/design_pdb.py
pengzhangzhi/ab_opt
from diffab.tools.runner.design_for_testset import main

if __name__ == '__main__':
    main()

AbDesign/design_testset.py
pengzhangzhi/ab_opt
from .sabdab import SAbDabDataset
from .custom import CustomDataset

from ._base import get_dataset

AbDesign/diffab/datasets/__init__.py
pengzhangzhi/ab_opt
from torch.utils.data import Dataset, ConcatDataset
from diffab.utils.transforms import get_transform


_DATASET_DICT = {}


def register_dataset(name):
    def decorator(cls):
        _DATASET_DICT[name] = cls
        return cls
    return decorator


def get_dataset(cfg):
    transform = get_transform(cfg.transform) if 'transform' in cfg else None
    return _DATASET_DICT[cfg.type](cfg, transform=transform)


@register_dataset('concat')
def get_concat_dataset(cfg):
    datasets = [get_dataset(d) for d in cfg.datasets]
    return ConcatDataset(datasets)


@register_dataset('balanced_concat')
class BalancedConcatDataset(Dataset):

    def __init__(self, cfg, transform=None):
        super().__init__()
        assert transform is None, 'transform is not supported.'
        self.datasets = [get_dataset(d) for d in cfg.datasets]
        self.max_size = max([len(d) for d in self.datasets])

    def __len__(self):
        return self.max_size * len(self.datasets)

    def __getitem__(self, idx):
        dataset_idx = idx // self.max_size
        return self.datasets[dataset_idx][idx % len(self.datasets[dataset_idx])]

AbDesign/diffab/datasets/_base.py
pengzhangzhi/ab_opt
import os
import logging
import joblib
import pickle
import lmdb
from Bio import PDB
from Bio.PDB import PDBExceptions
from torch.utils.data import Dataset
from tqdm.auto import tqdm

from ..utils.protein import parsers
from .sabdab import _label_heavy_chain_cdr, _label_light_chain_cdr
from ._base import register_dataset


def preprocess_antibody_structure(task):
    pdb_path = task['pdb_path']
    H_id = task.get('heavy_id', 'H')
    L_id = task.get('light_id', 'L')

    parser = PDB.PDBParser(QUIET=True)
    model = parser.get_structure(id, pdb_path)[0]

    all_chain_ids = [c.id for c in model]

    parsed = {
        'id': task['id'],
        'heavy': None,
        'heavy_seqmap': None,
        'light': None,
        'light_seqmap': None,
        'antigen': None,
        'antigen_seqmap': None,
    }
    try:
        if H_id in all_chain_ids:
            (
                parsed['heavy'], 
                parsed['heavy_seqmap']
            ) = _label_heavy_chain_cdr(*parsers.parse_biopython_structure(
                model[H_id],
                max_resseq = 113    # Chothia, end of Heavy chain Fv
            ))
        
        if L_id in all_chain_ids:
            (
                parsed['light'], 
                parsed['light_seqmap']
            ) = _label_light_chain_cdr(*parsers.parse_biopython_structure(
                model[L_id],
                max_resseq = 106    # Chothia, end of Light chain Fv
            ))

        if parsed['heavy'] is None and parsed['light'] is None:
            raise ValueError(
                f'Neither valid antibody H-chain or L-chain is found. '
                f'Please ensure that the chain id of heavy chain is "{H_id}" '
                f'and the id of the light chain is "{L_id}".'
            )

        
        ag_chain_ids = [cid for cid in all_chain_ids if cid not in (H_id, L_id)]
        if len(ag_chain_ids) > 0:
            chains = [model[c] for c in ag_chain_ids]
            (
                parsed['antigen'], 
                parsed['antigen_seqmap']
            ) = parsers.parse_biopython_structure(chains)

    except (
        PDBExceptions.PDBConstructionException, 
        parsers.ParsingException, 
        KeyError,
        ValueError,
    ) as e:
        logging.warning('[{}] {}: {}'.format(
            task['id'], 
            e.__class__.__name__, 
            str(e)
        ))
        return None

    return parsed


@register_dataset('custom')
class CustomDataset(Dataset):

    MAP_SIZE = 32*(1024*1024*1024)  # 32GB

    def __init__(self, structure_dir, transform=None, reset=False):
        super().__init__()
        self.structure_dir = structure_dir
        self.transform = transform

        self.db_conn = None
        self.db_ids = None
        self._load_structures(reset)

    @property
    def _cache_db_path(self):
        return os.path.join(self.structure_dir, 'structure_cache.lmdb')

    def _connect_db(self):
        self._close_db()
        self.db_conn = lmdb.open(
            self._cache_db_path,
            map_size=self.MAP_SIZE,
            create=False,
            subdir=False,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )
        with self.db_conn.begin() as txn:
            keys = [k.decode() for k in txn.cursor().iternext(values=False)]
            self.db_ids = keys

    def _close_db(self):
        if self.db_conn is not None:
            self.db_conn.close()
        self.db_conn = None
        self.db_ids = None
        
    def _load_structures(self, reset):
        all_pdbs = []
        for fname in os.listdir(self.structure_dir):
            if not fname.endswith('.pdb'): continue
            all_pdbs.append(fname)

        if reset or not os.path.exists(self._cache_db_path):
            todo_pdbs = all_pdbs
        else:
            self._connect_db()
            processed_pdbs = self.db_ids
            self._close_db()
            todo_pdbs = list(set(all_pdbs) - set(processed_pdbs))

        if len(todo_pdbs) > 0:
            self._preprocess_structures(todo_pdbs)
    
    def _preprocess_structures(self, pdb_list):
        tasks = []
        for pdb_fname in pdb_list:
            pdb_path = os.path.join(self.structure_dir, pdb_fname)
            tasks.append({
                'id': pdb_fname,
                'pdb_path': pdb_path,
            })

        data_list = joblib.Parallel(
            n_jobs = max(joblib.cpu_count() // 2, 1),
        )(
            joblib.delayed(preprocess_antibody_structure)(task)
            for task in tqdm(tasks, dynamic_ncols=True, desc='Preprocess')
        )

        db_conn = lmdb.open(
            self._cache_db_path,
            map_size = self.MAP_SIZE,
            create=True,
            subdir=False,
            readonly=False,
        )
        ids = []
        with db_conn.begin(write=True, buffers=True) as txn:
            for data in tqdm(data_list, dynamic_ncols=True, desc='Write to LMDB'):
                if data is None:
                    continue
                ids.append(data['id'])
                txn.put(data['id'].encode('utf-8'), pickle.dumps(data))

    def __len__(self):
        return len(self.db_ids)

    def __getitem__(self, index):
        self._connect_db()
        id = self.db_ids[index]
        with self.db_conn.begin() as txn:
            data = pickle.loads(txn.get(id.encode()))
        if self.transform is not None:
            data = self.transform(data)
        return data


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--dir', type=str, default='./data/custom')
    parser.add_argument('--reset', action='store_true', default=False)
    args = parser.parse_args()

    dataset = CustomDataset(
        structure_dir = args.dir,
        reset = args.reset,
    )
    print(dataset[0])
    print(len(dataset))
    
AbDesign/diffab/datasets/custom.py
pengzhangzhi/ab_opt
import os
import random
import logging
import datetime
import pandas as pd
import joblib
import pickle
import lmdb
import subprocess
import torch
from Bio import PDB, SeqRecord, SeqIO, Seq
from Bio.PDB import PDBExceptions
from Bio.PDB import Polypeptide
from torch.utils.data import Dataset
from tqdm.auto import tqdm

from diffab.utils.transforms.filter_structure import MissingChainException
from diffab.utils.transforms.patch import CUTOFFVIOLATION

from ..utils.protein import parsers, constants
from ._base import register_dataset


ALLOWED_AG_TYPES = {
    'protein',
    'protein | protein',
    'protein | protein | protein',
    'protein | protein | protein | protein | protein',
    'protein | protein | protein | protein',
}

RESOLUTION_THRESHOLD = 4.0

TEST_ANTIGENS = [
    'sars-cov-2 receptor binding domain',
    'hiv-1 envelope glycoprotein gp160',
    'mers s',
    'influenza a virus',
    'cd27 antigen',
]


def nan_to_empty_string(val):
    if val != val or not val:
        return ''
    else:
        return val


def nan_to_none(val):
    if val != val or not val:
        return None
    else:
        return val


def split_sabdab_delimited_str(val):
    if not val:
        return []
    else:
        return [s.strip() for s in val.split('|')]


def parse_sabdab_resolution(val):
    if val == 'NOT' or not val or val != val:
        return None
    elif isinstance(val, str) and ',' in val:
        return float(val.split(',')[0].strip())
    else:
        return float(val)


def _aa_tensor_to_sequence(aa):
    return ''.join([Polypeptide.index_to_one(a.item()) for a in aa.flatten()])


def _label_heavy_chain_cdr(data, seq_map, max_cdr3_length=30):
    if data is None or seq_map is None:
        return data, seq_map

    # Add CDR labels
    cdr_flag = torch.zeros_like(data['aa'])
    for position, idx in seq_map.items():
        resseq = position[1]
        cdr_type = constants.ChothiaCDRRange.to_cdr('H', resseq)
        if cdr_type is not None:
            cdr_flag[idx] = cdr_type
    data['cdr_flag'] = cdr_flag

    # Add CDR sequence annotations
    data['H1_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.H1] )
    data['H2_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.H2] )
    data['H3_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.H3] )

    cdr3_length = (cdr_flag == constants.CDR.H3).sum().item()
    # Remove too long CDR3
    if cdr3_length > max_cdr3_length:
        cdr_flag[cdr_flag == constants.CDR.H3] = 0
        logging.warning(f'CDR-H3 too long {cdr3_length}. Removed.')
        return None, None

    # Filter: ensure CDR3 exists
    if cdr3_length == 0:
        logging.warning('No CDR-H3 found in the heavy chain.')
        return None, None

    return data, seq_map


def _label_light_chain_cdr(data, seq_map, max_cdr3_length=30):
    if data is None or seq_map is None:
        return data, seq_map
    cdr_flag = torch.zeros_like(data['aa'])
    for position, idx in seq_map.items():
        resseq = position[1]
        cdr_type = constants.ChothiaCDRRange.to_cdr('L', resseq)
        if cdr_type is not None:
            cdr_flag[idx] = cdr_type
    data['cdr_flag'] = cdr_flag

    data['L1_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.L1] )
    data['L2_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.L2] )
    data['L3_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.L3] )

    cdr3_length = (cdr_flag == constants.CDR.L3).sum().item()
    # Remove too long CDR3
    if cdr3_length > max_cdr3_length:
        cdr_flag[cdr_flag == constants.CDR.L3] = 0
        logging.warning(f'CDR-L3 too long {cdr3_length}. Removed.')
        return None, None

    # Ensure CDR3 exists
    if cdr3_length == 0:
        logging.warning('No CDRs found in the light chain.')
        return None, None

    return data, seq_map


def preprocess_sabdab_structure(task):
    try:
        entry = task['entry']
        pdb_path = task['pdb_path']

        parser = PDB.PDBParser(QUIET=True)
        model = parser.get_structure(id, pdb_path)[0]

        parsed = {
            'id': entry['id'],
            'heavy': None,
            'heavy_seqmap': None,
            'light': None,
            'light_seqmap': None,
            'antigen': None,
            'antigen_seqmap': None,
        }
        if entry['H_chain'] is not None:
            (
                parsed['heavy'], 
                parsed['heavy_seqmap']
            ) = _label_heavy_chain_cdr(*parsers.parse_biopython_structure(
                model[entry['H_chain']],
                max_resseq = 113    # Chothia, end of Heavy chain Fv
            ))
        
        if entry['L_chain'] is not None:
            (
                parsed['light'], 
                parsed['light_seqmap']
            ) = _label_light_chain_cdr(*parsers.parse_biopython_structure(
                model[entry['L_chain']],
                max_resseq = 106    # Chothia, end of Light chain Fv
            ))

        if parsed['heavy'] is None and parsed['light'] is None:
            raise ValueError('Neither valid H-chain or L-chain is found.')
    
        if len(entry['ag_chains']) > 0:
            chains = [model[c] for c in entry['ag_chains']]
            (
                parsed['antigen'], 
                parsed['antigen_seqmap']
            ) = parsers.parse_biopython_structure(chains)

    except (
        PDBExceptions.PDBConstructionException, 
        parsers.ParsingException, 
        KeyError,
        ValueError,
    ) as e:
        logging.warning('[{}] {}: {}'.format(
            task['id'], 
            e.__class__.__name__, 
            str(e)
        ))
        return None

    return parsed


class SAbDabDataset(Dataset):

    MAP_SIZE = 32*(1024*1024*1024)  # 32GB

    def __init__(
        self, 
        summary_path = './data/sabdab_summary_all.tsv', 
        chothia_dir = './data/all_structures/chothia', 
        processed_dir = './data/processed',
        split = 'train',
        split_seed = 2022,
        transform = None,
        reset = False,
    ):
        super().__init__()
        self.summary_path = summary_path
        self.chothia_dir = chothia_dir
        if not os.path.exists(chothia_dir):
            raise FileNotFoundError(
                f"SAbDab structures not found in {chothia_dir}. "
                "Please download them from http://opig.stats.ox.ac.uk/webapps/newsabdab/sabdab/archive/all/"
            )
        self.processed_dir = processed_dir
        os.makedirs(processed_dir, exist_ok=True)

        self.sabdab_entries = None
        self._load_sabdab_entries()

        self.db_conn = None
        self.db_ids = None
        self._load_structures(reset)

        self.clusters = None
        self.id_to_cluster = None
        self._load_clusters(reset)

        self.ids_in_split = None
        self._load_split(split, split_seed)

        self.transform = transform

    def _load_sabdab_entries(self):
        df = pd.read_csv(self.summary_path, sep='\t')
        entries_all = []
        for i, row in tqdm(
            df.iterrows(), 
            dynamic_ncols=True, 
            desc='Loading entries',
            total=len(df),
        ):
            entry_id = "{pdbcode}_{H}_{L}_{Ag}".format(
                pdbcode = row['pdb'],
                H = nan_to_empty_string(row['Hchain']),
                L = nan_to_empty_string(row['Lchain']),
                Ag = ''.join(split_sabdab_delimited_str(
                    nan_to_empty_string(row['antigen_chain'])
                ))
            )
            ag_chains = split_sabdab_delimited_str(
                nan_to_empty_string(row['antigen_chain'])
            )
            resolution = parse_sabdab_resolution(row['resolution'])
            entry = {
                'id': entry_id,
                'pdbcode': row['pdb'],
                'H_chain': nan_to_none(row['Hchain']),
                'L_chain': nan_to_none(row['Lchain']),
                'ag_chains': ag_chains,
                'ag_type': nan_to_none(row['antigen_type']),
                'ag_name': nan_to_none(row['antigen_name']),
                'date': datetime.datetime.strptime(row['date'], '%m/%d/%y'),
                'resolution': resolution,
                'method': row['method'],
                'scfv': row['scfv'],
            }

            # Filtering
            if (
                (entry['ag_type'] in ALLOWED_AG_TYPES or entry['ag_type'] is None)
                and (entry['resolution'] is not None and entry['resolution'] <= RESOLUTION_THRESHOLD)
            ):
                entries_all.append(entry)
        self.sabdab_entries = entries_all

    def _load_structures(self, reset):
        if not os.path.exists(self._structure_cache_path) or reset:
            if os.path.exists(self._structure_cache_path):
                os.unlink(self._structure_cache_path)
            self._preprocess_structures()

        with open(self._structure_cache_path + '-ids', 'rb') as f:
            self.db_ids = pickle.load(f)
        self.sabdab_entries = list(
            filter(
                lambda e: e['id'] in self.db_ids,
                self.sabdab_entries
            )
        )

    @property
    def _structure_cache_path(self):
        return os.path.join(self.processed_dir, 'structures.lmdb')
        
    def _preprocess_structures(self):
        tasks = []
        for entry in self.sabdab_entries:
            pdb_path = os.path.join(self.chothia_dir, '{}.pdb'.format(entry['pdbcode']))
            if not os.path.exists(pdb_path):
                logging.warning(f"PDB not found: {pdb_path}")
                continue
            tasks.append({
                'id': entry['id'],
                'entry': entry,
                'pdb_path': pdb_path,
            })

        data_list = joblib.Parallel(
            n_jobs = max(joblib.cpu_count() // 2, 1),
        )(
            joblib.delayed(preprocess_sabdab_structure)(task)
            for task in tqdm(tasks, dynamic_ncols=True, desc='Preprocess')
        )

        db_conn = lmdb.open(
            self._structure_cache_path,
            map_size = self.MAP_SIZE,
            create=True,
            subdir=False,
            readonly=False,
        )
        ids = []
        with db_conn.begin(write=True, buffers=True) as txn:
            for data in tqdm(data_list, dynamic_ncols=True, desc='Write to LMDB'):
                if data is None:
                    continue
                ids.append(data['id'])
                txn.put(data['id'].encode('utf-8'), pickle.dumps(data))

        with open(self._structure_cache_path + '-ids', 'wb') as f:
            pickle.dump(ids, f)

    @property
    def _cluster_path(self):
        return os.path.join(self.processed_dir, 'cluster_result_cluster.tsv')

    def _load_clusters(self, reset):
        if not os.path.exists(self._cluster_path) or reset:
            self._create_clusters()

        clusters, id_to_cluster = {}, {}
        with open(self._cluster_path, 'r') as f:
            for line in f.readlines():
                cluster_name, data_id = line.split()
                if cluster_name not in clusters:
                    clusters[cluster_name] = []
                clusters[cluster_name].append(data_id)
                id_to_cluster[data_id] = cluster_name
        self.clusters = clusters
        self.id_to_cluster = id_to_cluster

    def _create_clusters(self):
        cdr_records = []
        for id in self.db_ids:
            structure = self.get_structure(id)
            if structure['heavy'] is not None:
                cdr_records.append(SeqRecord.SeqRecord(
                    Seq.Seq(structure['heavy']['H3_seq']),
                    id = structure['id'],
                    name = '',
                    description = '',
                ))
            elif structure['light'] is not None:
                cdr_records.append(SeqRecord.SeqRecord(
                    Seq.Seq(structure['light']['L3_seq']),
                    id = structure['id'],
                    name = '',
                    description = '',
                ))
        fasta_path = os.path.join(self.processed_dir, 'cdr_sequences.fasta')
        SeqIO.write(cdr_records, fasta_path, 'fasta')

        cmd = ' '.join([
            'mmseqs', 'easy-cluster',
            os.path.realpath(fasta_path),
            'cluster_result', 'cluster_tmp',
            '--min-seq-id', '0.5',
            '-c', '0.8',
            '--cov-mode', '1',
        ])
        subprocess.run(cmd, cwd=self.processed_dir, shell=True, check=True)

    def _load_split(self, split, split_seed):
        assert split in ('train', 'val', 'test')
        ids_test = [
            entry['id']
            for entry in self.sabdab_entries
            if entry['ag_name'] in TEST_ANTIGENS
        ]
        test_relevant_clusters = set([self.id_to_cluster[id] for id in ids_test])

        ids_train_val = [
            entry['id']
            for entry in self.sabdab_entries
            if self.id_to_cluster[entry['id']] not in test_relevant_clusters
        ]
        random.Random(split_seed).shuffle(ids_train_val)
        if split == 'test':
            self.ids_in_split = ids_test
        elif split == 'val':
            self.ids_in_split = ids_train_val[:20]
        else:
            self.ids_in_split = ids_train_val[20:]

    def _connect_db(self):
        if self.db_conn is not None:
            return
        self.db_conn = lmdb.open(
            self._structure_cache_path,
            map_size=self.MAP_SIZE,
            create=False,
            subdir=False,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )

    def get_structure(self, id):
        self._connect_db()
        with self.db_conn.begin() as txn:
            return pickle.loads(txn.get(id.encode()))

    def __len__(self):
        return len(self.ids_in_split)

    def __getitem__(self, index):
        id = self.ids_in_split[index]
        data = self.get_structure(id)
        if self.transform is not None:
            try:
                data = self.transform(data)
            except (MissingChainException,CUTOFFVIOLATION) as e:
                return self[random.randint(0, len(self) - 1)]
        return data


@register_dataset('sabdab')
def get_sabdab_dataset(cfg, transform):
    return SAbDabDataset(
        summary_path = cfg.summary_path,
        chothia_dir = cfg.chothia_dir,
        processed_dir = cfg.processed_dir,
        split = cfg.split,
        split_seed = cfg.get('split_seed', 2022),
        transform = transform,
    )


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--split', type=str, default='train')
    parser.add_argument('--processed_dir', type=str, default='./data/processed')
    parser.add_argument('--reset', action='store_true', default=False)
    args = parser.parse_args()
    if args.reset:
        sure = input('Sure to reset? (y/n): ')
        if sure != 'y':
            exit()
    dataset = SAbDabDataset(
        processed_dir=args.processed_dir,
        split=args.split, 
        reset=args.reset
    )
    print(dataset[0])
    print(len(dataset), len(dataset.clusters))

AbDesign/diffab/datasets/sabdab.py
pengzhangzhi/ab_opt
from .diffab import DiffusionAntibodyDesign
from .diff_gnn import DiffusionGNN
from ._base import get_model

AbDesign/diffab/models/__init__.py
pengzhangzhi/ab_opt

_MODEL_DICT = {}


def register_model(name):
    def decorator(cls):
        _MODEL_DICT[name] = cls
        return cls
    return decorator


def get_model(cfg):
    if cfg.type == 'diff_gnn':
        return _MODEL_DICT[cfg.type](**cfg)
    else:
         return _MODEL_DICT[cfg.type](cfg)

AbDesign/diffab/models/_base.py
pengzhangzhi/ab_opt
import functools
from random import random
from torch_scatter import scatter_sum
import torch
import torch.nn as nn
from diffab.modules.MCAttGNN.mc_att_model import ProteinFeature
from diffab.modules.MCAttGNN.mc_egnn import MCAttEGNN, SeqGNN
from diffab.modules.diffusion.transition import PositionTransition, RotationTransition
from diffab.utils.protein.constants import num_aa_types,AA
from diffab.modules.common.geometry import construct_3d_basis, local_to_global
from diffab.modules.common.so3 import random_uniform_so3, rotation_to_so3vec, so3vec_to_rotation
from diffab.modules.encoders.residue import ResidueEmbedding
from diffab.modules.encoders.pair import PairEmbedding
from diffab.modules.diffusion.dpm_full import rotation_matrix_cosine_loss
from diffab.utils.protein.constants import max_num_heavyatoms, BBHeavyAtom,backbone_atom_coordinates_tensor
from ._base import register_model
import torch.nn.functional as F
from tqdm import tqdm
from diffab.utils.misc import pair2edge,batchfy,clash_loss
resolution_to_num_atoms = {
    'backbone+CB': 5,
    'full': max_num_heavyatoms
}


@register_model('diff_gnn')
class DiffusionGNN(nn.Module):

    def __init__(
        self,
        res_feat_dim=128,
        edge_feat_dim=1,
        hidden_size=128,
        num_steps=100,
        num_atoms=3,
        n_layers=6,
        dropout=0.1, 
        dense=False,
        self_condition=False,
        hotspot=True,
        objective='pred_x0', 
        trans_rot_opt={}, 
        trans_pos_opt={}, 
        position_mean=[0.0],
        position_scale=[10.0],
        train_structure=True,
        train_sequence=True,
        *args, **kwargs
        ):
        super().__init__()

        self.num_atoms = 3
        self.num_steps = num_steps
        self.self_condition = self_condition
        assert objective in ['pred_x0', 'pred_noise']
        self.objective = objective
        self.protein_edge_constructor = ProteinFeature()
        self.embed = GraphEmbedding(
            res_feat_dim, edge_feat_dim, self.num_atoms,
            hotspot=hotspot,remove_structure=train_structure,
            remove_sequence=train_sequence
            )

        self.gnn = MCAttEGNN(
            res_feat_dim, hidden_size, hidden_size,
            num_atoms, edge_feat_dim, n_layers=n_layers,
            residual=True, dropout=dropout, dense=dense,
            self_condition=self_condition
        )
        self.seqgnn = SeqGNN(
            hidden_size, hidden_size, num_aa_types,
            num_atoms, 0
            )

        self.trans_rot = RotationTransition(num_steps, **trans_rot_opt)
        self.trans_pos = PositionTransition(num_steps, **trans_pos_opt)
        self.register_buffer('position_mean', torch.FloatTensor(position_mean))
        self.register_buffer('position_scale', torch.FloatTensor(position_scale))
        self.register_buffer('_dummy', torch.empty([0, ]))
    
    def _normalize_position(self, p):
        p_norm = (p - self.position_mean) / self.position_scale
        return p_norm

    def _unnormalize_position(self, p_norm):
        p = p_norm * self.position_scale + self.position_mean
        return p
        
    def init_mask(self, aa, coord, generate_range):
        """
        mask the generated residues of aa and coord.
        construct pseudo coordinates of `to-be-generated` residues by interpolating them from the neighboring residues.
        this prevents from leaking the information of `to-be-generated` residues
        Args:
            aa (torch.Tensor): [nb_all_res,]
            coord (torch.Tensor): [nb_all_res, nb_atoms, 3]
            generate_range (torch.Tensor): [bs, 2]
        """
        coord, aa, cmask = coord.clone(), aa.clone(), torch.zeros_like(coord)
        n_channel, n_dim = coord.shape[1:]
        for start, end in generate_range:
            aa[start:end + 1] = AA.UNK
            l_coord, r_coord = coord[start - 1], coord[end + 1]  # [n_channel, 3]
            n_span = end - start + 2
            coord_offsets = (r_coord - l_coord).unsqueeze(0).expand(n_span - 1, n_channel, n_dim)  # [n_mask, n_channel, 3]
            coord_offsets = torch.cumsum(coord_offsets, dim=0)
            mask_coords = l_coord + coord_offsets / n_span
            coord[start:end + 1] = mask_coords
            cmask[start:end + 1, ...] = 1
        return coord, aa, cmask
    
    def forward(self, batch):
        graph_data = batch['graph']
        coord = graph_data['coord'][...,:self.num_atoms,:]
        aa = graph_data['aa']
        batch_id = graph_data['batch_id']
        masked_coord, masked_aa, cmask = self.init_mask(aa, coord, graph_data['genearte_flag_range'])
        mask_gen = cmask[:, 0, 0].bool()  # [n_all_node]
        ctx_edges, inter_edges, ctx_edge_feats =self.protein_edge_constructor(
            masked_coord, batch_id, graph_data['fragment_id']
        )
        lengths = scatter_sum(torch.ones_like(batch_id), batch_id)
        h0, ctx_edge_feats, iter_edge_feats = self.embed(batch, ctx_edges,inter_edges, lengths)
        coord = self._normalize_position(coord)
        R_0 = construct_3d_basis(
            coord[...,BBHeavyAtom.CA,:],
            coord[...,BBHeavyAtom.C,:],
            coord[...,BBHeavyAtom.N,:],
        )
        v_0 = rotation_to_so3vec(R_0)
        p_0 = coord[...,BBHeavyAtom.CA,:]


        bs = lengths.shape[0]
        t = torch.randint(0, self.num_steps, (bs,), dtype=torch.long, device=self._dummy.device)
        # [nb_all_res, ]
        expanded_t = torch.cat([time.repeat(l) for time,l in zip(t,lengths)])

        v_noisy, _ = self.trans_rot.add_noise(v_0, mask_gen, expanded_t)
        p_noisy, eps_p = self.trans_pos.add_noise(p_0, mask_gen,  expanded_t)

        noised_coord = reconstruct_noised_coord(coord, p_noisy, v_noisy, mask_gen)
        noise_condition = self.trans_pos.var_sched.betas[expanded_t]
        h_self_cond = None
        hotspot_label = batch['hotspot_label'][batch['mask']] # [nb_all_res,]
        # original hotspot_label contains 2 for hotspot and 1 for non-hotspot. 
        # cast to 1 for hotspot and 0 for non-hotspot.
        hotspot_label = torch.where(hotspot_label == 2, torch.ones_like(hotspot_label), torch.zeros_like(hotspot_label))
        hydropathy = batch['hydropathy'][batch['mask']]
        charge = batch['charge'][batch['mask']]
        h, z = self.gnn(
            h0, noised_coord,
            ctx_edges, inter_edges,
            noise_condition,class_condition=hotspot_label, 
            hydropathy=hydropathy, charge=charge,
            h_self_cond=h_self_cond,
            ctx_edge_attr=ctx_edge_feats,
            att_edge_attr=None,
        )
        aa_logits = self.seqgnn(
            h, inter_edges,z,None,
            hydropathy=hydropathy, charge=charge,
            )
        R_pred = construct_3d_basis(
            z[...,BBHeavyAtom.CA,:],
            z[...,BBHeavyAtom.C,:],
            z[...,BBHeavyAtom.N,:],
        )
        p_pred = z[...,BBHeavyAtom.CA,:]
        if self.objective == 'pred_x0':
            p0_pred = p_pred
            eps_p_pred = self.trans_pos.pred_noise_from_start(p_noisy, p0_pred, mask_gen, expanded_t)
            p_true = p_0
        elif self.objective == 'pred_noise':
            NotImplementedError()
        loss_dict = {}
        loss_rot = rotation_matrix_cosine_loss(R_pred[mask_gen], R_0[mask_gen]).mean()
        loss_dict['rot'] = loss_rot

        loss_pos = F.mse_loss(p_pred[mask_gen], p_true[mask_gen], reduction='none').sum(dim=-1).mean()
        loss_dict['Ca-pos'] = loss_pos 
        
        bb_loss = F.mse_loss(z[mask_gen], coord[mask_gen], reduction='none').sum(dim=-1).mean()
        loss_dict['bb-pos'] = bb_loss
        # seq loss
        loss_seq = F.cross_entropy(aa_logits[mask_gen], aa[mask_gen])
        loss_dict['seq'] = loss_seq 
        accuracy = torch.sum(aa_logits[mask_gen].softmax(dim=-1).argmax(dim=-1) == aa[mask_gen]) / aa[mask_gen].size(0)
        loss_dict['accuracy'] = accuracy.item() 
        
        
        # clash loss
        # Experiment shows that
        # 1) clash loss converges very soon. 
        # 2) result in RMSD increase by 0.5A in test set. 
        # 3) predicted CDR overly extend to the epitope region 
        # batch_p0_pred = batchfy(p0_pred,lengths)
        # loss_clash = clash_loss(batch_p0_pred,batch['mask'],batch['chain_nb'],3.0078/self.position_scale)
        # loss_dict['cdr2epitope_clash'] = loss_clash
        return loss_dict

    @torch.no_grad()
    def sample(
        self, 
        batch, 
        sample_opt={
            'sample_structure': True,
            'sample_sequence': True,
        }
    ):
        graph_data = batch['graph']
        coord = graph_data['coord'][...,:self.num_atoms,:]
        aa = graph_data['aa']
        batch_id = graph_data['batch_id']
        masked_coord, masked_aa, cmask = self.init_mask(aa, coord, graph_data['genearte_flag_range'])
        mask_gen = cmask[:, 0, 0].bool()  # [n_all_node]
        ctx_edges, inter_edges, ctx_edge_feats =self.protein_edge_constructor(
            masked_coord, batch_id, graph_data['fragment_id']
        )
        coord = self._normalize_position(coord)
        R_0 = construct_3d_basis(
            coord[...,BBHeavyAtom.CA,:],
            coord[...,BBHeavyAtom.C,:],
            coord[...,BBHeavyAtom.N,:],
        )
        v_0 = rotation_to_so3vec(R_0)
        p_0 = coord[...,BBHeavyAtom.CA,:]

        lengths = scatter_sum(torch.ones_like(batch_id), batch_id)
        bs = lengths.shape[0]
        nb_all_res = v_0.shape[0]
        h0, ctx_edge_feats, iter_edge_feats = self.embed(batch, ctx_edges,inter_edges, lengths)
        
        if sample_opt['sample_structure']:
            v_rand = random_uniform_so3([nb_all_res], device=self._dummy.device)
            p_rand = torch.randn_like(p_0) # make_interpolate_positions(p, mask_generate) # 
            v_init = torch.where(mask_gen[..., None].expand_as(v_0), v_rand, v_0)
            p_init = torch.where(mask_gen[..., None].expand_as(p_0), p_rand, p_0)
        else:
            v_init, p_init = v_0, p_0
        s_init = aa
        
        # [nb_all_res,]
        traj = {self.num_steps: (v_init, self._unnormalize_position(p_init), s_init)}
        pbar = functools.partial(tqdm, total=self.num_steps, desc='Sampling')
        start = None
        for t in pbar(range(self.num_steps, 0, -1)):
            v_t, p_t, s_t = traj[t]
            p_t = self._normalize_position(p_t)
            
            t_tensor = torch.full((bs,), fill_value=t, dtype=torch.long, device=self._dummy.device)
            # [nb_all_res, ]
            expanded_t = torch.cat([time.repeat(l) for time,l in zip(t_tensor,lengths)])
            noise_condition = self.trans_pos.var_sched.betas[expanded_t]
         
            noised_coord = reconstruct_noised_coord(coord, p_t, v_t, mask_gen)
            self_cond = start if self.self_condition else None
            hotspot_label = batch['hotspot_label'][batch['mask']] # [nb_all_res,]
            # original hotspot_label contains 2 for hotspot and 1 for non-hotspot. 
            # cast to 1 for hotspot and 0 for non-hotspot.
            hotspot_label = torch.where(hotspot_label == 2, torch.ones_like(hotspot_label), torch.zeros_like(hotspot_label))
            hydropathy = batch['hydropathy'][batch['mask']] 
            charge = batch['charge'][batch['mask']] 
            h, z = self.gnn(
                h0, noised_coord,
                ctx_edges, inter_edges,
                noise_condition,class_condition=hotspot_label, 
                hydropathy=hydropathy, charge=charge,
                h_self_cond=self_cond,
                ctx_edge_attr=ctx_edge_feats,
                att_edge_attr=None,
            )
            aa_logits = self.seqgnn(
            h, inter_edges,z,None,
            hydropathy=hydropathy, charge=charge,
            )
            R_pred = construct_3d_basis(
                z[...,BBHeavyAtom.CA,:],
                z[...,BBHeavyAtom.C,:],
                z[...,BBHeavyAtom.N,:],
            )
            v_next = rotation_to_so3vec(R_pred)
            p_pred = z[...,BBHeavyAtom.CA,:]
            if self.objective == 'pred_x0':
                p0_pred = p_pred
                eps_p_pred = self.trans_pos.pred_noise_from_start(p_t, p0_pred, mask_gen, expanded_t)
            elif self.objective == 'pred_noise':
                NotImplementedError()
            
            v_next = self.trans_rot.denoise(v_t, v_next, mask_gen, expanded_t)
            # p_next = self.trans_pos.denoise(p_t, eps_p, mask_gen, expanded_t)
            p_next = self.trans_pos.denoise_from_p0(p_t, p0_pred, mask_gen, expanded_t,guidance_kwargs={**batch,"lengths":lengths})
            
            s_next = aa_logits.softmax(dim=-1).argmax(dim=-1)
            traj[t-1] = (v_next, self._unnormalize_position(p_next), s_next)
            traj[t] = tuple(x.cpu() for x in traj[t])    # Move previous states to cpu memory.
        
        for k,v in traj.items():
            # batchfy the trajectory data
            # to be compatible with subsequent code
            traj[k] = tuple(batchfy(x,lengths) for x in v)
        
        return traj
    @torch.no_grad()
    def optimize(
        self, 
        batch, 
        opt_step, 
        optimize_opt={
            'sample_structure': True,
            'sample_sequence': True,
        }
    ):
        mask_generate = batch['generate_flag']
        mask_res = batch['mask']
        res_feat, pair_feat, R_0, p_0 = self.encode(
            batch,
            remove_structure = optimize_opt.get('sample_structure', True),
            remove_sequence = optimize_opt.get('sample_sequence', True)
        )
        v_0 = rotation_to_so3vec(R_0)
        s_0 = batch['aa']

        traj = self.diffusion.optimize(v_0, p_0, s_0, opt_step, res_feat, pair_feat, mask_generate, mask_res, **optimize_opt)
        return traj

class GraphEmbedding(nn.Module):
    def __init__(self, res_dim, pair_dim, max_num_atoms=3,hotspot=True,remove_structure=True, remove_sequence=True):
        super().__init__()
        self.max_num_atoms = max_num_atoms
        self.max_aa_types = num_aa_types
        self.remove_structure = remove_structure
        self.remove_sequence = remove_sequence
        self.residue_embed = ResidueEmbedding(res_dim, max_num_atoms, hotspot=hotspot)
        self.pair_embed = PairEmbedding(pair_dim, max_num_atoms)

    def forward(self, batch, ctx_edges, iter_edges, lengths):
        # info of generated regions is not leaked out
        mask = torch.logical_and(
            batch['mask_heavyatom'][:, :, BBHeavyAtom.CA], 
            ~batch['generate_flag']     # Context means ``not generated''
        )
        structure_mask = mask if self.remove_structure else None
        sequence_mask = mask if self.remove_sequence else None
        # [B, N, D]
        res_feat = self.residue_embed(
            aa = batch['aa'],
            res_nb = batch['res_nb'],
            chain_nb = batch['chain_nb'],
            pos_atoms = batch['pos_heavyatom'],
            mask_atoms = batch['mask_heavyatom'],
            fragment_type = batch['fragment_type'],
            hotspot =  batch['hotspot_label'],
            # hydropathy = batch['hydropathy'],
            # charge = batch['charge'],
            structure_mask = structure_mask,
            sequence_mask = sequence_mask,
        )
        pair_feat = self.pair_embed(
            aa = batch['aa'],
            res_nb = batch['res_nb'],
            chain_nb = batch['chain_nb'],
            pos_atoms = batch['pos_heavyatom'],
            mask_atoms = batch['mask_heavyatom'],
            structure_mask = structure_mask,
            sequence_mask = sequence_mask,
        ) # (B, L, L, D)

        pad_mask = batch['mask'] # (B, N)
        node_feat = res_feat[pad_mask]
        ctx_edge_feats = pair2edge(ctx_edges, lengths, pair_feat)
        iter_edge_feats = pair2edge(iter_edges, lengths, pair_feat)
        return node_feat, ctx_edge_feats, iter_edge_feats

def reconstruct_noised_coord(coord, p_noisy, v_noisy, mask_gen):
    """
    Args:
        coord: (N, num_atoms, 3)
        p_noisy: (N, num_atoms, 3)
        v_noisy: (N, num_atoms, 3)
        mask_gen: (N, ) bool, true denotes the residue to be generated
    """
    coord = coord.clone()
    num_reconstructed_res = mask_gen.sum()
    device = coord.device
    r_noisy = so3vec_to_rotation(v_noisy[mask_gen])
    bb_coords = backbone_atom_coordinates_tensor.clone().to(device)  # (21, 3, 3)
    # [num_reconstructed_res, 3, 3]
    bb_coords = bb_coords[0].unsqueeze(0).repeat(num_reconstructed_res, 1, 1)
    reconstructed_bb_pos = local_to_global(r_noisy, p_noisy[mask_gen], bb_coords) 
    coord[mask_gen] = reconstructed_bb_pos
    return coord


def pairwise_distance_loss(x_true,x_pred,mask):

    true_dist = torch.cdist(x_true,x_true)
    pred_dist = torch.cdist(x_pred,x_pred)
    mask_dist = torch.cdist(mask,mask)
    loss = torch.mean((true_dist-pred_dist).abs()*mask_dist)
    return loss
AbDesign/diffab/models/diff_gnn.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn

from diffab.modules.common.geometry import construct_3d_basis
from diffab.modules.common.so3 import rotation_to_so3vec
from diffab.modules.encoders.residue import ResidueEmbedding
from diffab.modules.encoders.pair import PairEmbedding
from diffab.modules.diffusion.dpm_full import FullDPM
from diffab.utils.protein.constants import max_num_heavyatoms, BBHeavyAtom
from ._base import register_model


resolution_to_num_atoms = {
    'backbone+CB': 5,
    'full': max_num_heavyatoms
}


@register_model('diffab')
class DiffusionAntibodyDesign(nn.Module):

    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg

        num_atoms = resolution_to_num_atoms[cfg.get('resolution', 'full')]
        self.residue_embed = ResidueEmbedding(cfg.res_feat_dim, num_atoms)
        self.pair_embed = PairEmbedding(cfg.pair_feat_dim, num_atoms)

        self.diffusion = FullDPM(
            cfg.res_feat_dim,
            cfg.pair_feat_dim,
            **cfg.diffusion,
        )

    def encode(self, batch, remove_structure, remove_sequence):
        """
        Returns:
            res_feat:   (N, L, res_feat_dim)
            pair_feat:  (N, L, L, pair_feat_dim)
        """
        # This is used throughout embedding and encoding layers
        #   to avoid data leakage.
        context_mask = torch.logical_and(
            batch['mask_heavyatom'][:, :, BBHeavyAtom.CA], 
            ~batch['generate_flag']     # Context means ``not generated''
        )

        structure_mask = context_mask if remove_structure else None
        sequence_mask = context_mask if remove_sequence else None

        res_feat = self.residue_embed(
            aa = batch['aa'],
            res_nb = batch['res_nb'],
            chain_nb = batch['chain_nb'],
            pos_atoms = batch['pos_heavyatom'],
            mask_atoms = batch['mask_heavyatom'],
            fragment_type = batch['fragment_type'],
            structure_mask = structure_mask,
            sequence_mask = sequence_mask,
        )

        pair_feat = self.pair_embed(
            aa = batch['aa'],
            res_nb = batch['res_nb'],
            chain_nb = batch['chain_nb'],
            pos_atoms = batch['pos_heavyatom'],
            mask_atoms = batch['mask_heavyatom'],
            structure_mask = structure_mask,
            sequence_mask = sequence_mask,
        )

        R = construct_3d_basis(
            batch['pos_heavyatom'][:, :, BBHeavyAtom.CA],
            batch['pos_heavyatom'][:, :, BBHeavyAtom.C],
            batch['pos_heavyatom'][:, :, BBHeavyAtom.N],
        )
        p = batch['pos_heavyatom'][:, :, BBHeavyAtom.CA]

        return res_feat, pair_feat, R, p
    
    def forward(self, batch):
        mask_generate = batch['generate_flag']
        mask_res = batch['mask']
        res_feat, pair_feat, R_0, p_0 = self.encode(
            batch,
            remove_structure = self.cfg.get('train_structure', True),
            remove_sequence = self.cfg.get('train_sequence', True)
        )
        v_0 = rotation_to_so3vec(R_0)
        s_0 = batch['aa']

        loss_dict = self.diffusion(
            v_0, p_0, s_0, res_feat, pair_feat, mask_generate, mask_res,
            denoise_structure = self.cfg.get('train_structure', True),
            denoise_sequence  = self.cfg.get('train_sequence', True),
        )
        return loss_dict

    @torch.no_grad()
    def sample(
        self, 
        batch, 
        sample_opt={
            'sample_structure': True,
            'sample_sequence': True,
        }
    ):
        mask_generate = batch['generate_flag']
        mask_res = batch['mask']
        res_feat, pair_feat, R_0, p_0 = self.encode(
            batch,
            remove_structure = sample_opt.get('sample_structure', True),
            remove_sequence = sample_opt.get('sample_sequence', True)
        )
        v_0 = rotation_to_so3vec(R_0)
        s_0 = batch['aa']
        traj = self.diffusion.sample(v_0, p_0, s_0, res_feat, pair_feat, mask_generate, mask_res, **sample_opt)
        return traj

    @torch.no_grad()
    def optimize(
        self, 
        batch, 
        opt_step, 
        optimize_opt={
            'sample_structure': True,
            'sample_sequence': True,
        }
    ):
        mask_generate = batch['generate_flag']
        mask_res = batch['mask']
        res_feat, pair_feat, R_0, p_0 = self.encode(
            batch,
            remove_structure = optimize_opt.get('sample_structure', True),
            remove_sequence = optimize_opt.get('sample_sequence', True)
        )
        v_0 = rotation_to_so3vec(R_0)
        s_0 = batch['aa']

        traj = self.diffusion.optimize(v_0, p_0, s_0, opt_step, res_feat, pair_feat, mask_generate, mask_res, **optimize_opt)
        return traj

AbDesign/diffab/models/diffab.py
pengzhangzhi/ab_opt
#!/usr/bin/python
# -*- coding:utf-8 -*-
from .mc_att_model import EfficientMCAttModel

AbDesign/diffab/modules/MCAttGNN/__init__.py
pengzhangzhi/ab_opt
#!/usr/bin/python
# -*- coding:utf-8 -*-
import torch
from torch import nn
import torch.nn.functional as F

from torch_scatter import scatter_sum

from .pdb_utils import VOCAB

from .mc_egnn import MCAttEGNN


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


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


class ProteinFeature(nn.Module):

    def __init__(self):
        super().__init__()
        ...

    @torch.no_grad()
    def construct_edges(self, X, batch_id, segment_ids):
        '''
        Memory efficient with complexity of O(Nn) where n is the largest number of nodes in the batch
        '''
        # construct tensors to map between global / local node index
        lengths = scatter_sum(torch.ones_like(batch_id), batch_id)  # [bs]
        offsets = F.pad(torch.cumsum(lengths, dim=0)[:-1], pad=(1, 0), value=0)  # [bs]
        N, max_n = batch_id.shape[0], torch.max(lengths)
        # global node index to local index. lni2gni can be implemented as lni + offsets[batch_id]
        gni = torch.arange(N, device=batch_id.device)
        gni2lni = gni - offsets[batch_id]  # [N]


        ctx_edges, inter_edges = [], []

        # all possible edges (within the same graph)
        # same bid (get rid of self-loop and none edges)
        same_bid = torch.zeros(N, max_n, device=batch_id.device)
        same_bid[(gni, lengths[batch_id] - 1)] = 1
        same_bid = 1 - torch.cumsum(same_bid, dim=-1)
        # shift right and pad 1 to the left
        same_bid = F.pad(same_bid[:, :-1], pad=(1, 0), value=1)
        same_bid[(gni, gni2lni)] = 0  # delete self loop
        row, col = torch.nonzero(same_bid).T  # [2, n_edge_all]
        col = col + offsets[batch_id[row]]  # mapping from local to global node index
        
        # not global edges
        # is_global = sequential_or(S == self.boa_idx, S == self.boh_idx, S == self.bol_idx) # [N]
        # row_global, col_global = is_global[row], is_global[col]
        # not_global_edges = torch.logical_not(torch.logical_or(row_global, col_global))

        # all possible ctx edges: same seg, not global
        row_seg, col_seg = segment_ids[row], segment_ids[col]
        select_edges = (row_seg == col_seg) # the src and dst node are in the same segment
        ctx_all_row, ctx_all_col = row[select_edges], col[select_edges]
        # ctx edges
        ctx_edges = _radial_edges(X, torch.stack([ctx_all_row, ctx_all_col]).T, cutoff=8.0)

        # all possible inter edges: not same seg, not global
        select_edges = (row_seg != col_seg)
        inter_all_row, inter_all_col = row[select_edges], col[select_edges]
        inter_edges = _radial_edges(X, torch.stack([inter_all_row, inter_all_col]).T, cutoff=12.0)

        num_ctx_edge = ctx_edges.shape[-1]
        ctx_edge_feats = torch.zeros(num_ctx_edge, 1, device=X.device) # [num_ctx_edge, 1]
        
        return ctx_edges, inter_edges, ctx_edge_feats

    def forward(self, X, batch_id, segment_ids):
        """
        

        Args:
            X (torch.Tensor): [nb_all_residues, nb_atoms, 3]
            batch_id (torch.Tensor): [nb_all_residues], which sample the residue belongs to
            segment_ids (torch.Tensor): [nb_all_residues], which segment the residue belongs to
            cdr_range (torch.Tensor): [bs, 2]

        Returns:
            ctx_edges, edges within the same segment
            inter_edges, egdes between different segments
            ctx_edge_feats, [num_ctx_edge, 1] initialized to 0
        """
        
        ctx_edges, inter_edges, ctx_edge_feats = self.construct_edges(X,batch_id,segment_ids)
        return ctx_edges, inter_edges, ctx_edge_feats

def _radial_edges(X, src_dst, cutoff):
    dist = X[:, 1][src_dst]  # [Ef, 2, 3], CA position
    dist = torch.norm(dist[:, 0] - dist[:, 1], dim=-1) # [Ef]
    src_dst = src_dst[dist <= cutoff]
    src_dst = src_dst.transpose(0, 1)  # [2, Ef]
    return src_dst


class MCAttModel(nn.Module):
    def __init__(self, embed_size, hidden_size, n_channel, n_edge_feats=0,
                 n_layers=3, cdr_type='3', alpha=0.1, dropout=0.1, dense=False):
        super().__init__()
        self.num_aa_type = len(VOCAB)
        self.cdr_type = cdr_type
        self.mask_token_id = VOCAB.get_unk_idx()
        self.alpha = alpha

        self.aa_embedding = nn.Embedding(self.num_aa_type, embed_size)
        self.gnn = MCAttEGNN(embed_size, hidden_size, self.num_aa_type,
                             n_channel, n_edge_feats, n_layers=n_layers,
                             residual=True, dropout=dropout, dense=dense)
        
        # self.seq_loss = nn.CrossEntropyLoss(reduction='none')
        # self.coord_loss = nn.MSELoss(reduction='sum')
        # it is the same as HuberLoss with delta=1.0
        # self.coord_loss = nn.SmoothL1Loss(reduction='sum')
        
        # edges
        self.protein_feature = ProteinFeature()

    def seq_loss(self, _input, target):
        return F.cross_entropy(_input, target, reduction='none')

    def coord_loss(self, _input, target):
        return F.smooth_l1_loss(_input, target, reduction='sum')

    def init_mask(self, X, S, cdr_range):
        '''
        set coordinates of masks following a unified distribution
        between the two ends
        '''
        X, S, cmask = X.clone(), S.clone(), torch.zeros_like(X, device=X.device)
        n_channel, n_dim = X.shape[1:]
        for start, end in cdr_range:
            S[start:end + 1] = self.mask_token_id
            l_coord, r_coord = X[start - 1], X[end + 1]  # [n_channel, 3]
            n_span = end - start + 2
            coord_offsets = (r_coord - l_coord).unsqueeze(0).expand(n_span - 1, n_channel, n_dim)  # [n_mask, n_channel, 3]
            coord_offsets = torch.cumsum(coord_offsets, dim=0)
            mask_coords = l_coord + coord_offsets / n_span
            X[start:end + 1] = mask_coords
            cmask[start:end + 1, ...] = 1
        return X, S, cmask
    

    def forward(self, X, S, L, offsets):
        '''
        :param X: [n_all_node, n_channel, 3]
        :param S: [n_all_node]
        :param L: list of cdr types
        :param offsets: [batch_size + 1]
        '''
        # prepare inputs
        cdr_range = torch.tensor(
            [(cdr.index(self.cdr_type), cdr.rindex(self.cdr_type)) for cdr in L],
            dtype=torch.long, device=X.device
        ) + offsets[:-1].unsqueeze(-1)
        max_step = max([r - l + 1 for l, r in cdr_range])
        pred_positions, aa_cnt = [], 0
        for t in range(max_step):
            position = cdr_range[:, 0] + t
            position = position.masked_select(position <= cdr_range[:, 1])
            pred_positions.append(position)
            aa_cnt += len(position)

        # save ground truth
        true_X, true_S = X.clone(), S.clone()

        # init mask
        X, S, cmask = self.init_mask(X, S, cdr_range)

        snll, closs = 0, 0
        for t in range(max_step):
            with torch.no_grad():
                ctx_edges, inter_edges, ctx_edge_feats = self.protein_feature(X, S, offsets)

            pred_position = pred_positions[t]

            H_0 = self.aa_embedding(S)  # [n_all_node, embed_size]
            H, Z = self.gnn(H_0, X, ctx_edges, inter_edges, ctx_edge_feats)
            
            # loss
            logits = H[pred_position]
            gt_s, gt_x = true_S[pred_position], true_X[pred_position]
            snll = snll + torch.sum(self.seq_loss(logits, gt_s)) / aa_cnt
            # closs = closs + self.coord_loss(Z[pred_position], gt_x) / aa_cnt
            cur_cmask = torch.zeros_like(cmask, device=cmask.device, dtype=torch.bool)
            cur_cmask[pred_position] = 1
            closs = closs + self.coord_loss(Z.masked_select(cur_cmask), true_X.masked_select(cur_cmask)) / aa_cnt
            
            # update X, S
            S = S.clone()
            S[pred_position] = true_S[pred_position]
            # pure autoregressive
            # cmask = cmask.clone()
            # cmask[pred_position] = 0  # use ground truth
            X = (1 - cmask) * true_X + cmask * Z

        loss = snll + self.alpha * closs
        return loss, snll, closs


    def generate(self, X, S, L, offsets, greedy=True):
        '''
        :param X: [n_all_node, n_channel, 3]
        :param S: [n_all_node]
        :param L: list of cdr types
        :param offsets: [batch_size + 1]
        '''
        # prepare inputs
        cdr_range = torch.tensor(
            [(cdr.index(self.cdr_type), cdr.rindex(self.cdr_type)) for cdr in L],
            dtype=torch.long, device=X.device
        ) + offsets[:-1].unsqueeze(-1)
        max_step = max([r - l + 1 for l, r in cdr_range])
        pred_positions, aa_cnt = [], 0
        for t in range(max_step):
            position = cdr_range[:, 0] + t
            position = position.masked_select(position <= cdr_range[:, 1])
            pred_positions.append(position)
            aa_cnt += len(position)

        # save ground truth
        true_X, true_S = X.clone(), S.clone()

        # init mask
        X, S, cmask = self.init_mask(X, S, cdr_range)

        special_mask = torch.tensor(VOCAB.get_special_mask(), device=S.device, dtype=torch.long)

        snll_all_step = []
        for t in range(max_step):
            with torch.no_grad():
                ctx_edges, inter_edges, ctx_edge_feats = self.protein_feature(X, S, offsets)

            pred_position = pred_positions[t]

            H_0 = self.aa_embedding(S)  # [n_all_node, embed_size]
            H, Z = self.gnn(H_0, X, ctx_edges, inter_edges, ctx_edge_feats)
            
            # sequence
            smask = special_mask.repeat(len(pred_position), 1).bool()  # [n, vocab_size]
            logits = H[pred_position]
            logits = logits.masked_fill(smask, float('-inf'))  # mask special tokens
            if greedy:
                S[pred_position] = torch.argmax(logits, dim=-1)  # [n]
            else:
                prob = F.softmax(logits, dim=-1)  # [n, vocab_size]
                S[pred_position] = torch.multinomial(prob, num_samples=1).squeeze()

            # nll loss
            snll_all_step.append(self.seq_loss(logits, S[pred_position]))
            
            # update X
            # cmask = cmask.clone()
            # cmask[pred_position] = 0  # fix already predicted X
            X = (1 - cmask) * X + cmask * Z

        return snll_all_step, S, X, true_X, cdr_range

    # unified api: return ppl (list), seq (list), x (numpy), true_x (numpy) and whether aligned
    def infer(self, batch, device):
        X, S, L, offsets = batch['X'].to(device), batch['S'].to(device), batch['L'], batch['offsets'].to(device)
        snll_all_step, pred_S, pred_X, true_X, cdr_range = self.generate(
            X, S, L, offsets
        )
        pred_S, cdr_range = pred_S.tolist(), cdr_range.tolist()
        pred_X, true_X = pred_X.cpu().numpy(), true_X.cpu().numpy()
        # seqs, x, true_x
        seq, x, true_x = [], [], []
        for start, end in cdr_range:
            end = end + 1
            seq.append(''.join([VOCAB.idx_to_symbol(pred_S[i]) for i in range(start, end)]))
            x.append(pred_X[start:end])
            true_x.append(true_X[start:end])
        # ppl
        ppl = [0 for _ in range(len(cdr_range))]
        lens = [0 for _ in ppl]
        for t, snlls in enumerate(snll_all_step):
            i, j = 0, 0
            if len(snlls.shape) == 0:
                snlls = snlls.reshape(1)
            for start, end in cdr_range:
                if start + t <= end:
                    ppl[i] += snlls[j]
                    lens[i] += 1
                    j += 1
                i += 1
        ppl = [p / n for p, n in zip(ppl, lens)]
        ppl = torch.exp(torch.tensor(ppl, device=device)).tolist()
        return ppl, seq, x, true_x, True


class EfficientMCAttModel(MCAttModel):
    def __init__(self, embed_size, hidden_size, n_channel, n_edge_feats=0, n_layers=3, cdr_type='3', alpha=0.1, dropout=0.1, n_iter=5):
        super().__init__(embed_size, hidden_size, n_channel, n_edge_feats, n_layers, cdr_type, alpha, dropout, dense=False)
        self.n_iter = n_iter

    def forward(self, X, S, L, offsets):
        '''
        :param X: [n_all_node, n_channel, 3]
        :param S: [n_all_node]
        :param L: list of cdr types
        :param offsets: [batch_size + 1]
        '''
        # prepare inputs
        cdr_range = torch.tensor(
            [(cdr.index(self.cdr_type), cdr.rindex(self.cdr_type)) for cdr in L],
            dtype=torch.long, device=X.device
        ) + offsets[:-1].unsqueeze(-1)

        # save ground truth
        true_X, true_S = X.clone(), S.clone()

        # init mask
        # mask in this code context means that the invalid positions denoted by 1.
        X, S, cmask = self.init_mask(X, S, cdr_range)  # [n_all_node, n_channel, 3]
        mask = cmask[:, 0, 0].bool()  # [n_all_node]
        aa_cnt = mask.sum()

        special_mask = torch.tensor(VOCAB.get_special_mask(), device=S.device, dtype=torch.long)
        smask = special_mask.repeat(aa_cnt, 1).bool()
        H_0 = self.aa_embedding(S)  # [n_all_node, embed_size]
        aa_embeddings = self.aa_embedding(torch.arange(self.num_aa_type, device=H_0.device))  # [vocab_size, embed_size]

        snll, closs = 0, 0

        for r in range(self.n_iter):
            with torch.no_grad():
                ctx_edges, inter_edges, ctx_edge_feats = self.protein_feature(X, S, offsets)
            H, Z = self.gnn(H_0, X, ctx_edges, inter_edges, ctx_edge_feats)

            # refine
            X = X.clone()
            X[mask] = Z[mask]
            H_0 = H_0.clone()
            logits = H[mask]
            seq_prob = torch.softmax(logits.masked_fill(smask, float('-inf')), dim=-1)  # [aa_cnt, vocab_size]
            H_0[mask] = seq_prob.mm(aa_embeddings)  # smooth embedding
            
            r_snll = torch.sum(self.seq_loss(logits, true_S[mask])) / aa_cnt
            snll += r_snll / self.n_iter
        
        closs = self.coord_loss(Z[mask], true_X[mask]) / aa_cnt

        loss = snll + self.alpha * closs
        return loss, r_snll, closs  # only return the last snll

    def generate(self, X, S, L, offsets, greedy=True):
        '''
        :param X: [n_all_node, n_channel, 3]
        :param S: [n_all_node]
        :param L: list of cdr types
        :param offsets: [batch_size + 1]
        '''
        # prepare inputs
        cdr_range = torch.tensor(
            [(cdr.index(self.cdr_type), cdr.rindex(self.cdr_type)) for cdr in L],
            dtype=torch.long, device=X.device
        ) + offsets[:-1].unsqueeze(-1)

        # save ground truth
        true_X, true_S = X.clone(), S.clone()

        # init mask
        X, S, cmask = self.init_mask(X, S, cdr_range)
        mask = cmask[:, 0, 0].bool()  # [n_all_node]
        aa_cnt = mask.sum()

        special_mask = torch.tensor(VOCAB.get_special_mask(), device=S.device, dtype=torch.long)
        smask = special_mask.repeat(aa_cnt, 1).bool()
        H_0 = self.aa_embedding(S)  # [n_all_node, embed_size]
        aa_embeddings = self.aa_embedding(torch.arange(self.num_aa_type, device=H_0.device))  # [vocab_size, embed_size]
        
        for r in range(self.n_iter):
            with torch.no_grad():
                ctx_edges, inter_edges, ctx_edge_feats = self.protein_feature(X, S, offsets)
            H, Z = self.gnn(H_0, X, ctx_edges, inter_edges, ctx_edge_feats)

            # refine
            X = X.clone()
            X[mask] = Z[mask]
            H_0 = H_0.clone()
            seq_prob = torch.softmax(H[mask].masked_fill(smask, float('-inf')), dim=-1)  # [aa_cnt, vocab_size]
            H_0[mask] = seq_prob.mm(aa_embeddings)  # smooth embedding

        # H_0 = self.aa_embedding(S)  # [n_all_node, embed_size]
        # H, Z = self.gnn(H_0, X, ctx_edges, inter_edges, ctx_edge_feats)

        logits = H[mask]  # [aa_cnt, vocab_size]

        logits = logits.masked_fill(smask, float('-inf'))  # mask special tokens

        if greedy:
            S[mask] = torch.argmax(logits, dim=-1)  # [n]
        else:
            prob = F.softmax(logits, dim=-1)
            S[mask] = torch.multinomial(prob, num_samples=1).squeeze()
        snll_all = self.seq_loss(logits, S[mask])

        return snll_all, S, X, true_X, cdr_range

    def infer(self, batch, device, greedy=True):
        X, S, L, offsets = batch['X'].to(device), batch['S'].to(device), batch['L'], batch['offsets'].to(device)
        snll_all, pred_S, pred_X, true_X, cdr_range = self.generate(
            X, S, L, offsets, greedy=greedy
        )
        pred_S, cdr_range = pred_S.tolist(), cdr_range.tolist()
        pred_X, true_X = pred_X.cpu().numpy(), true_X.cpu().numpy()
        # seqs, x, true_x
        seq, x, true_x = [], [], []
        for start, end in cdr_range:
            end = end + 1
            seq.append(''.join([VOCAB.idx_to_symbol(pred_S[i]) for i in range(start, end)]))
            x.append(pred_X[start:end])
            true_x.append(true_X[start:end])
        # ppl
        ppl = [0 for _ in range(len(cdr_range))]
        lens = [0 for _ in ppl]
        offset = 0
        for i, (start, end) in enumerate(cdr_range):
            length = end - start + 1
            for t in range(length):
                ppl[i] += snll_all[t + offset]
            offset += length
            lens[i] = length
        ppl = [p / n for p, n in zip(ppl, lens)]
        ppl = torch.exp(torch.tensor(ppl, device=device)).tolist()
        return ppl, seq, x, true_x, True

    def generate_analyze(self, X, S, L, offsets, greedy=True):
        '''
        :param X: [n_all_node, n_channel, 3]
        :param S: [n_all_node]
        :param L: list of cdr types
        :param offsets: [batch_size + 1]
        '''
        # prepare inputs
        cdr_range = torch.tensor(
            [(cdr.index(self.cdr_type), cdr.rindex(self.cdr_type)) for cdr in L],
            dtype=torch.long, device=X.device
        ) + offsets[:-1].unsqueeze(-1)

        # save ground truth
        true_X, true_S = X.clone(), S.clone()

        # init mask
        X, S, cmask = self.init_mask(X, S, cdr_range)
        mask = cmask[:, 0, 0].bool()  # [n_all_node]
        aa_cnt = mask.sum()

        special_mask = torch.tensor(VOCAB.get_special_mask(), device=S.device, dtype=torch.long)
        smask = special_mask.repeat(aa_cnt, 1).bool()
        H_0 = self.aa_embedding(S)  # [n_all_node, embed_size]
        aa_embeddings = self.aa_embedding(torch.arange(self.num_aa_type, device=H_0.device))  # [vocab_size, embed_size]
        
        round_probs, round_Xs = [], []

        for r in range(self.n_iter):
            with torch.no_grad():
                ctx_edges, inter_edges, ctx_edge_feats = self.protein_feature(X, S, offsets)
            H, Z, atts = self.gnn(H_0, X, ctx_edges, inter_edges, ctx_edge_feats, return_attention=True)

            # refine
            X = X.clone()
            X[mask] = Z[mask]
            H_0 = H_0.clone()
            masked_logits = H[mask].masked_fill(smask, float('-inf'))
            seq_prob = torch.softmax(masked_logits, dim=-1)  # [aa_cnt, vocab_size]
            H_0[mask] = seq_prob.mm(aa_embeddings)  # smooth embedding

            round_probs.append(masked_logits)
            round_Xs.append(X)

        # H_0 = self.aa_embedding(S)  # [n_all_node, embed_size]
        # H, Z = self.gnn(H_0, X, ctx_edges, inter_edges, ctx_edge_feats)

        logits = H[mask]  # [aa_cnt, vocab_size]

        logits = logits.masked_fill(smask, float('-inf'))  # mask special tokens

        if greedy:
            S[mask] = torch.argmax(logits, dim=-1)  # [n]
        else:
            prob = F.softmax(logits, dim=-1)
            S[mask] = torch.multinomial(prob, num_samples=1).squeeze()  # [n]
        snll_all = self.seq_loss(logits, S[mask])

        return {
            'r_probs': round_probs,
            'r_Xs': round_Xs,
            'cdr_range': cdr_range,
            'final_atts': atts
        }



def interpolate(X, cdr_range):
    '''
    set coordinates of masks following a unified distribution
    between the two ends
    '''
    X, cmask = X.clone(),  torch.zeros_like(X, device=X.device)
    n_channel, n_dim = X.shape[1:]
    for start, end in cdr_range:
        l_coord, r_coord = X[start - 1], X[end + 1]  # [n_channel, 3]
        n_span = end - start + 2
        coord_offsets = (r_coord - l_coord).unsqueeze(0).expand(n_span - 1, n_channel, n_dim)  # [n_mask, n_channel, 3]
        coord_offsets = torch.cumsum(coord_offsets, dim=0)
        mask_coords = l_coord + coord_offsets / n_span
        X[start:end + 1] = mask_coords
        cmask[start:end + 1, ...] = 1
    return X, cmask
AbDesign/diffab/modules/MCAttGNN/mc_att_model.py
pengzhangzhi/ab_opt
#!/usr/bin/python
# -*- coding:utf-8 -*-
'''
    Most codes are copied from https://github.com/vgsatorras/egnn, which is the official implementation of
    the paper:
        E(n) Equivariant Graph Neural Networks
        Victor Garcia Satorras, Emiel Hogeboom, Max Welling
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch_scatter import scatter_softmax
from diffab.modules.common.geometry import quaternion_1ijk_to_rotation_matrix

from diffab.modules.common.so3 import so3vec_to_rotation


def exists(x):
    return x is not None
class MC_E_GCL(nn.Module):
    """
    Multi-Channel E(n) Equivariant Convolutional Layer
    """

    def __init__(self, input_nf, output_nf, hidden_nf, n_channel, edges_in_d=0, act_fn=nn.SiLU(),
                 residual=True, attention=False, normalize=False, coords_agg='mean', tanh=False,
                 dropout=0.1,time_emb_dim = None,):
        super(MC_E_GCL, self).__init__()
        input_edge = input_nf * 2
        self.residual = residual
        self.attention = attention
        self.normalize = normalize
        self.coords_agg = coords_agg
        self.tanh = tanh
        self.epsilon = 1e-8

        self.dropout = nn.Dropout(dropout)

        self.edge_mlp = nn.Sequential(
            nn.Linear(input_edge + 2*(n_channel)**2 + edges_in_d, hidden_nf),
            act_fn,
            nn.Linear(hidden_nf, hidden_nf),
            act_fn)

        self.node_mlp = nn.Sequential(
            nn.Linear(hidden_nf + input_nf, hidden_nf),
            act_fn,
            nn.Linear(hidden_nf, output_nf))

        layer = nn.Linear(hidden_nf, n_channel, bias=False)
        torch.nn.init.xavier_uniform_(layer.weight, gain=0.001)

        coord_mlp = []
        coord_mlp.append(nn.Linear(hidden_nf, hidden_nf))
        coord_mlp.append(act_fn)
        coord_mlp.append(layer)
        if self.tanh:
            coord_mlp.append(nn.Tanh())
        self.coord_mlp = nn.Sequential(*coord_mlp)

        
        # rot_layer = nn.Linear(hidden_nf, n_channel, bias=False)
        # torch.nn.init.xavier_uniform_(rot_layer.weight, gain=0.001)
        # rot_mlp = []
        # rot_mlp.append(nn.Linear(hidden_nf, hidden_nf))
        # rot_mlp.append(act_fn)
        # rot_mlp.append(rot_layer)
        # if self.tanh:
        #     rot_mlp.append(nn.Tanh())
        # self.rot_mlp = nn.Sequential(*rot_mlp)
        
        # self.to_rot = nn.Linear(n_channel * 3, 3, bias=True)
        # with torch.no_grad():
        #     self.to_rot.weight.fill_(0.0)
            
        if self.attention:
            self.att_mlp = nn.Sequential(
                nn.Linear(hidden_nf, 1),
                nn.Sigmoid())
        self.time_mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(time_emb_dim, output_nf * 2)
        ) if exists(time_emb_dim) else None
        
    def edge_model(self, source, target, radial, edge_attr):
        '''
        :param source: [n_edge, input_size]
        :param target: [n_edge, input_size]
        :param radial: [n_edge, n_channel, n_channel]
        :param edge_attr: [n_edge, edge_dim]
        '''
        radial = radial.reshape(radial.shape[0], -1)  # [n_edge, n_channel ^ 2]

        if edge_attr is None:  # Unused.
            out = torch.cat([source, target, radial], dim=1)
        else:
            out = torch.cat([source, target, radial, edge_attr], dim=1)
        out = self.edge_mlp(out)
        out = self.dropout(out)

        if self.attention:
            att_val = self.att_mlp(out)
            out = out * att_val
        return out

    def node_model(self, x, edge_index, edge_attr, node_attr):
        '''
        :param x: [bs * n_node, input_size]
        :param edge_index: list of [n_edge], [n_edge]
        :param edge_attr: [n_edge, hidden_size], refers to message from i to j
        :param node_attr: [bs * n_node, node_dim]
        '''
        row, col = edge_index
        agg = unsorted_segment_sum(edge_attr, row, num_segments=x.size(0))  # [bs * n_node, hidden_size]
        # print(f'agg1, {torch.isnan(agg).sum()}', )
        if node_attr is not None:
            agg = torch.cat([x, agg, node_attr], dim=1)
        else:
            agg = torch.cat([x, agg], dim=1)  # [bs * n_node, input_size + hidden_size]
        # print(f'agg, {torch.isnan(agg).sum()}', )
        out = self.node_mlp(agg)  # [bs * n_node, output_size]
        # print(f'out, {torch.isnan(out).sum()}', )
        out = self.dropout(out)
        if self.residual:
            out = x + out
        return out, agg

    def coord_model(self, coord, edge_index, coord_diff, edge_feat):
        '''
        coord: [bs * n_node, n_channel, d]
        edge_index: list of [n_edge], [n_edge]
        coord_diff: [n_edge, n_channel, d]
        edge_feat: [n_edge, hidden_size]
        '''
        row, col = edge_index
        trans = coord_diff * self.coord_mlp(edge_feat).unsqueeze(-1)  # [n_edge, n_channel, d]

        if self.coords_agg == 'sum':
            agg = unsorted_segment_sum(trans, row, num_segments=coord.size(0))
        elif self.coords_agg == 'mean':
            agg = unsorted_segment_mean(trans, row, num_segments=coord.size(0))  # [bs * n_node, n_channel, d]
        else:
            raise Exception('Wrong coords_agg parameter' % self.coords_agg)
        coord = coord + agg
        return coord
    
    # def rot_model(self, rot, edge_index, coord_diff, edge_feat):
    #     '''
    #     rot: [bs * n_node, 3, 3]
    #     edge_index: list of [n_edge], [n_edge]
    #     coord_diff: [n_edge, n_channel, d]
    #     edge_feat: [n_edge, hidden_size]
    #     '''
    #     row, col = edge_index
    #     msg = coord_diff * self.rot_mlp(edge_feat).unsqueeze(-1)  # [n_edge, n_channel, d]

    #     if self.coords_agg == 'sum':
    #         agg = unsorted_segment_sum(msg, row, num_segments=rot.size(0))
    #     elif self.coords_agg == 'mean':
    #         agg = unsorted_segment_mean(msg, row, num_segments=rot.size(0))  # [bs * n_node, n_channel, d]
    #     else:
    #         raise Exception('Wrong coords_agg parameter' % self.coords_agg)
    #     n_nodes = rot.shape[0]
    #     agg = agg.reshape(n_nodes, -1) # [bs * n_node, n_channel * d]
    #     rot_update_vec = self.to_rot(agg) # [bs * n_node, 3]
    #     rot = rot @ quaternion_1ijk_to_rotation_matrix(rot_update_vec)
    #     # quat = rot_to_quat(rot) # [bs * n_node, 4]
    #     # rot = compose_q_update_vec(quat, rot_update_vec) # [bs * n_node, 3, 3]
    #     return rot

    def forward(self, h, edge_index, coord, edge_attr=None, node_attr=None,time_emb=None):
        '''
        h: [bs * n_node, hidden_size]
        edge_index: list of [n_row] and [n_col] where n_row == n_col (with no cutoff, n_row == bs * n_node * (n_node - 1))
        coord: [bs * n_node, n_channel, d]
        '''
        row, col = edge_index
        radial, dist, coord_diff = coord2radial(edge_index, coord)
        # concatenate radial and dist at the last two dimensions
        radial = torch.cat([radial, dist], dim=-1) # [n_edge, n_channel, 2 * n_channel]
        edge_feat = self.edge_model(h[row], h[col], radial, edge_attr)  # [n_edge, hidden_size]
        # print(f'edge_feat, {torch.isnan(edge_feat).sum()}', )
        coord = self.coord_model(coord, edge_index, coord_diff, edge_feat)    # [bs * n_node, n_channel, d]
        h, agg = self.node_model(h, edge_index, edge_feat, node_attr)
        if exists(self.time_mlp) and exists(time_emb):
            time_emb = self.time_mlp(time_emb)
            scale, shift = time_emb.chunk(2, dim = 1)
            h = h * (scale + 1) + shift

        return h, coord


class MC_Att_L(nn.Module):
    """
    Multi-Channel Attention Layer
    """
    def __init__(self, input_nf, output_nf, hidden_nf, n_channel, edges_in_d=0,
                 act_fn=nn.SiLU(), dropout=0.1, time_emb_dim = None,**kwargs):
        super().__init__()
        self.hidden_nf = hidden_nf

        self.dropout = nn.Dropout(dropout)

        self.linear_q = nn.Linear(input_nf, hidden_nf)
        self.linear_kv = nn.Linear(input_nf + 2*n_channel ** 2 + edges_in_d, hidden_nf * 2)  # parallel calculate kv

        layer = nn.Linear(hidden_nf, n_channel, bias=False)
        torch.nn.init.xavier_uniform_(layer.weight, gain=0.001)

        coord_mlp = []
        coord_mlp.append(nn.Linear(hidden_nf, hidden_nf))
        coord_mlp.append(act_fn)
        coord_mlp.append(layer)
        self.coord_mlp = nn.Sequential(*coord_mlp)
        self.time_mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(time_emb_dim, hidden_nf * 2)
        ) if exists(time_emb_dim) else None
        
        # rot_layer = nn.Linear(hidden_nf, n_channel, bias=False)
        # torch.nn.init.xavier_uniform_(rot_layer.weight, gain=0.001)
        # rot_mlp = []
        # rot_mlp.append(nn.Linear(hidden_nf, hidden_nf))
        # rot_mlp.append(act_fn)
        # rot_mlp.append(rot_layer)
        # self.rot_mlp = nn.Sequential(*rot_mlp)
        
        # self.to_rot = nn.Linear(n_channel * 3, 3, bias=True)
        # with torch.no_grad():
        #     self.to_rot.weight.fill_(0.0)
            
    def att_model(self, h, edge_index, radial, edge_attr):
        '''
        :param h: [bs * n_node, input_size]
        :param edge_index: list of [n_edge], [n_edge]
        :param radial: [n_edge, n_channel, n_channel]
        :param edge_attr: [n_edge, edge_dim]
        '''
        row, col = edge_index
        source, target = h[row], h[col]  # [n_edge, input_size]

        # qkv
        q = self.linear_q(source)  # [n_edge, hidden_size]
        n_channel = radial.shape[1]
        radial = radial.reshape(radial.shape[0], -1)  # [n_edge, n_channel ^ 2]
        if edge_attr is not None:
            target_feat = torch.cat([radial, target, edge_attr], dim=1)
        else:
            target_feat = torch.cat([radial, target], dim=1)
        kv = self.linear_kv(target_feat)  # [n_edge, hidden_size * 2]
        k, v = kv[..., 0::2], kv[..., 1::2]  # [n_edge, hidden_size]

        # attention weight
        alpha = torch.sum(q * k, dim=1)  # [n_edge]

        # print(f'alpha1, {torch.isnan(alpha).sum()}', )

        # alpha = scatter_softmax(alpha, row, h.shape[0]) # [n_edge]
        alpha = scatter_softmax(alpha, row) # [n_edge]

        # print(f'alpha2, {torch.isnan(alpha).sum()}', )

        return alpha, v
        
    def node_model(self, h, edge_index, att_weight, v):
        '''
        :param h: [bs * n_node, input_size]
        :param edge_index: list of [n_edge], [n_edge]
        :param att_weight: [n_edge, 1], unsqueezed before passed in
        :param v: [n_edge, hidden_size]
        '''
        row, _ = edge_index
        agg = unsorted_segment_sum(att_weight * v, row, h.shape[0])  # [bs * n_node, hidden_size]
        agg = self.dropout(agg)
        return h + agg  # residual

    def coord_model(self, coord, edge_index, coord_diff, att_weight, v):
        '''
        :param coord: [bs * n_node, n_channel, d]
        :param edge_index: list of [n_edge], [n_edge]
        :param coord_diff: [n_edge, n_channel, d]
        :param att_weight: [n_edge, 1], unsqueezed before passed in
        :param v: [n_edge, hidden_size]
        '''
        row, _ = edge_index
        coord_v = att_weight * self.coord_mlp(v)  # [n_edge, n_channel]
        trans = coord_diff * coord_v.unsqueeze(-1)
        agg = unsorted_segment_sum(trans, row, coord.size(0))
        coord = coord + agg
        return coord

    # def rot_model(self, rot, edge_index, coord_diff, att_weight, v):
    #     '''
    #     :param coord: [bs * n_node, n_channel, d]
    #     :param edge_index: list of [n_edge], [n_edge]
    #     :param coord_diff: [n_edge, n_channel, d]
    #     :param att_weight: [n_edge, 1], unsqueezed before passed in
    #     :param v: [n_edge, hidden_size]
    #     '''
    #     row, _ = edge_index
    #     coord_v = att_weight * self.rot_mlp(v)  # [n_edge, n_channel]
    #     trans = coord_diff * coord_v.unsqueeze(-1) # [n_edge, n_channel, d]
    #     agg = unsorted_segment_sum(trans, row, rot.size(0)) # [bs * n_node, n_channel, d]
    #     n_nodes = rot.shape[0]
    #     agg = agg.reshape(n_nodes, -1) # [bs * n_node, n_channel * d]
    #     rot_update_vec = self.to_rot(agg) # [bs * n_node, 3]
    #     quat = rot_to_quat(rot) # [bs * n_node, 4]
    #     rot = rot @ quaternion_1ijk_to_rotation_matrix(rot_update_vec)
    #     # rot = compose_q_update_vec(quat, rot_update_vec) # [bs * n_node, 3, 3]
    #     return rot
    
    def forward(self, h, edge_index, coord, edge_attr=None,time_emb=None):
        radial, dist, coord_diff = coord2radial(edge_index, coord)
        radial = torch.cat([radial, dist], dim=-1) # [n_edge, n_channel, n_channel * 2]
        att_weight, v = self.att_model(h, edge_index, radial, edge_attr)

        # print(f'att_weight, {torch.isnan(att_weight).sum()}', )
        # print(f'v, {torch.isnan(v).sum()}', )

        flat_att_weight = att_weight
        att_weight = att_weight.unsqueeze(-1)  # [n_edge, 1]
        h = self.node_model(h, edge_index, att_weight, v)
        coord = self.coord_model(coord, edge_index, coord_diff, att_weight, v)
        if exists(self.time_mlp) and exists(time_emb):
            time_emb = self.time_mlp(time_emb)
            scale, shift = time_emb.chunk(2, dim = 1)
            h = h * (scale + 1) + shift
        return h, coord, flat_att_weight


class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb
    

class SeqGNN(nn.Module):
    def __init__(self, in_node_nf, hidden_nf, out_node_nf, n_channel, in_edge_nf=0,
                 act_fn=nn.SiLU(), residual=True, dropout=0.1,
                 ):
        super().__init__()
        self.in_node_nf = in_node_nf
        self.dropout = nn.Dropout(dropout)
        


        self.linear_q = nn.Linear(in_node_nf, hidden_nf)
        self.linear_kv = nn.Linear(in_node_nf + 2*n_channel ** 2 + in_edge_nf, hidden_nf * 2)  # parallel calculate kv
        embed_dim = int(0.5*hidden_nf)
        self.hydropathy_embed = nn.Embedding(10, embed_dim, padding_idx=0)    # 1: hydrophilic, 2: moderate, 3. unknown
        self.charge_embed = nn.Embedding(10, embed_dim, padding_idx=0)    # 1: positive, 2: negative, 3. unknown
        fea_dim = embed_dim + embed_dim + hidden_nf # hydropathy + charge + h
        
        self.mlp = nn.Sequential(
            nn.Linear(fea_dim, fea_dim), 
            nn.ReLU(),
            nn.Linear(fea_dim, fea_dim),
            nn.ReLU(),
            nn.Linear(fea_dim, out_node_nf)
        )
        
        

    def att_model(self, h, edge_index, radial, edge_attr):
        '''
        :param h: [bs * n_node, input_size]
        :param edge_index: list of [n_edge], [n_edge]
        :param radial: [n_edge, n_channel, n_channel]
        :param edge_attr: [n_edge, edge_dim]
        '''
        row, col = edge_index
        source, target = h[row], h[col]  # [n_edge, input_size]

        # qkv
        q = self.linear_q(source)  # [n_edge, hidden_size]
        n_channel = radial.shape[1]
        radial = radial.reshape(radial.shape[0], -1)  # [n_edge, n_channel ^ 2]
        if edge_attr is not None:
            target_feat = torch.cat([radial, target, edge_attr], dim=1)
        else:
            target_feat = torch.cat([radial, target], dim=1)
        kv = self.linear_kv(target_feat)  # [n_edge, hidden_size * 2]
        k, v = kv[..., 0::2], kv[..., 1::2]  # [n_edge, hidden_size]

        # attention weight
        alpha = torch.sum(q * k, dim=1)  # [n_edge]

        # print(f'alpha1, {torch.isnan(alpha).sum()}', )

        # alpha = scatter_softmax(alpha, row, h.shape[0]) # [n_edge]
        alpha = scatter_softmax(alpha, row) # [n_edge]

        # print(f'alpha2, {torch.isnan(alpha).sum()}', )

        return alpha, v
        
    def node_model(self, h, edge_index, att_weight, v):
        '''
        :param h: [bs * n_node, input_size]
        :param edge_index: list of [n_edge], [n_edge]
        :param att_weight: [n_edge, 1], unsqueezed before passed in
        :param v: [n_edge, hidden_size]
        '''
        row, _ = edge_index
        agg = unsorted_segment_sum(att_weight * v, row, h.shape[0])  # [bs * n_node, hidden_size]
        agg = self.dropout(agg)
        return h + agg  # residual
    
    def forward(self, h, edge_index, coord, edge_attr, hydropathy, charge):
        radial, dist, coord_diff = coord2radial(edge_index, coord)
        radial = torch.cat([radial, dist], dim=-1)
        att_weight, v = self.att_model(h, edge_index, radial, edge_attr)
        att_weight = att_weight.unsqueeze(-1)  # [n_edge, 1]
        h = self.node_model(h, edge_index, att_weight, v)
        
        hydropathy_feat = self.hydropathy_embed(hydropathy)
        charge_feat = self.charge_embed(charge)

        h = torch.cat([h, hydropathy_feat, charge_feat], dim=-1)
        h = self.mlp(h)
        return h


class MCAttEGNN(nn.Module):
    def __init__(self, in_node_nf, hidden_nf, out_node_nf, n_channel, in_edge_nf=0,
                 act_fn=nn.SiLU(), n_layers=4, residual=True, dropout=0.1, dense=False,self_condition=False):
        super().__init__()
        '''
        :param in_node_nf: Number of features for 'h' at the input
        :param hidden_nf: Number of hidden features
        :param out_node_nf: Number of features for 'h' at the output
        :param n_channel: Number of channels of coordinates
        :param in_edge_nf: Number of features for the edge features
        :param act_fn: Non-linearity
        :param n_layers: Number of layer for the EGNN
        :param residual: Use residual connections, we recommend not changing this one
        :param dropout: probability of dropout
        :param dense: if dense, then context states will be concatenated for all layers,
                      coordination will be averaged
        '''
        self.hidden_nf = hidden_nf
        self.n_layers = n_layers
        self.self_condition = self_condition
        self.dropout = nn.Dropout(dropout)
        in_c = in_node_nf * 2 if self_condition else in_node_nf
        self.linear_in = nn.Linear(in_c, self.hidden_nf)

        
        for i in range(0, n_layers):
            self.add_module(f'gcl_{i}', MC_E_GCL(
                self.hidden_nf, self.hidden_nf, self.hidden_nf, n_channel,
                edges_in_d=in_edge_nf, act_fn=act_fn, residual=residual, dropout=dropout,time_emb_dim=self.hidden_nf,
            ))
            
            # TODO: add parameter for passing edge type to interaction layer
            self.add_module(f'att_{i}', MC_Att_L(
                self.hidden_nf, self.hidden_nf, self.hidden_nf,
                n_channel, edges_in_d=0, act_fn=act_fn, dropout=dropout,time_emb_dim=self.hidden_nf,
            ))
        self.out_layer = MC_E_GCL(
            self.hidden_nf, self.hidden_nf, self.hidden_nf, n_channel,
            edges_in_d=in_edge_nf, act_fn=act_fn, residual=residual,time_emb_dim=self.hidden_nf,
        )
        self.time_embedding = nn.Linear(self.hidden_nf + 3, self.hidden_nf)
        # embed_dim = int(0.5*self.hidden_nf)
        # self.hydropathy_embed = nn.Embedding(10, embed_dim, padding_idx=0)    # 1: hydrophilic, 2: moderate, 3. unknown
        # self.charge_embed = nn.Embedding(10, embed_dim, padding_idx=0)    # 1: positive, 2: negative, 3. unknown
        # fea_dim =  embed_dim + embed_dim + hidden_nf # hydropathy + charge + h
        # self.linear_embed = nn.Linear(fea_dim, hidden_nf)
        
        self.dense = dense
        if dense:
            self.linear_out = nn.Linear(self.hidden_nf * (n_layers + 1), out_node_nf)
        else:
            self.linear_out = nn.Linear(self.hidden_nf, out_node_nf)

        # self.rot_net = nn.Sequential(
        #     nn.Linear(self.hidden_nf, self.hidden_nf), nn.ReLU(),
        #     nn.Linear(self.hidden_nf, self.hidden_nf), nn.ReLU(),
        #     nn.Linear(self.hidden_nf, 3)
        # )
        # sinu_pos_emb = SinusoidalPosEmb(self.hidden_nf)

        # self.time_mlp = nn.Sequential(
        #     sinu_pos_emb,
        #     nn.Linear(self.hidden_nf, self.hidden_nf),
        #     nn.GELU(),
        #     nn.Linear(self.hidden_nf, self.hidden_nf)
        # )
        
        
    def forward(
        self, h, x, ctx_edges, att_edges, condition,  
        class_condition, hydropathy, charge,
        h_self_cond=None, ctx_edge_attr=None, att_edge_attr=None, 
        return_attention=False, return_h=False,
        ):
        if self.self_condition:
            h_self_cond = torch.zeros_like(h) if h_self_cond is None else h_self_cond
            h = torch.cat((h_self_cond, h), dim = -1)
        h = self.linear_in(h)
        h = self.dropout(h)
        
        # t = self.time_mlp(condition)
        
        ctx_states, ctx_coords, atts = [], [], []
        for i in range(0, self.n_layers):
            h, x = self._modules[f'gcl_{i}'](h, ctx_edges, x, edge_attr=ctx_edge_attr,time_emb=None)
            ctx_states.append(h)
            ctx_coords.append(x)
            h, x, att = self._modules[f'att_{i}'](h, att_edges, x, edge_attr=att_edge_attr,time_emb=None)
            atts.append(att)
            
        t_embed = torch.stack([condition, torch.sin(condition), torch.cos(condition)], dim=-1)
        h = self.time_embedding(torch.cat([h, t_embed], dim=-1))
        
        # # predict rotation
        # rot_vec_update = self.rot_net(h)
        # rot_mat_update = quaternion_1ijk_to_rotation_matrix(rot_vec_update)
        # rot = rot @ rot_mat_update
        
        # predict ca position
        h, x = self.out_layer(h, ctx_edges, x, edge_attr=ctx_edge_attr,time_emb=None)
        hidden = h
        ctx_states.append(h)
        ctx_coords.append(x)
        if self.dense:
            h = torch.cat(ctx_states, dim=-1)
            x = torch.mean(torch.stack(ctx_coords), dim=0)
            
        # predict aa
        # hydropathy_feat = self.hydropathy_embed(hydropathy)
        # charge_feat = self.charge_embed(charge)
        # h = torch.cat([h, hydropathy_feat, charge_feat], dim=-1)
        # h = self.linear_embed(h)
        h = self.dropout(h)
        h = self.linear_out(h) # [*, num_aa]
        out = (h, x)
        if return_attention:
            out = out + (atts,)
        if return_h:
            out = out + (hidden,)
        return out


def coord2radial(edge_index, coord):
    row, col = edge_index
    coord_diff = coord[row] - coord[col]  # [n_edge, n_channel, d]
    radial = torch.bmm(coord_diff, coord_diff.transpose(-1, -2))  # [n_edge, n_channel, n_channel]
    dist = torch.cdist(coord[row], coord[col])  # [n_edge, n_channel, n_channel]
    # normalize radial
    radial = F.normalize(radial, dim=0)  # [n_edge, n_channel, n_channel]
    dist = F.normalize(dist, dim=0)  # [n_edge, n_channel, n_channel]
    return radial, dist, coord_diff


def unsorted_segment_sum(data, segment_ids, num_segments):
    '''
    :param data: [n_edge, *dimensions]
    :param segment_ids: [n_edge]
    :param num_segments: [bs * n_node]
    '''
    expand_dims = tuple(data.shape[1:])
    result_shape = (num_segments, ) + expand_dims
    for _ in expand_dims:
        segment_ids = segment_ids.unsqueeze(-1)
    segment_ids = segment_ids.expand(-1, *expand_dims)
    result = data.new_full(result_shape, 0)  # Init empty result tensor.
    result.scatter_add_(0, segment_ids, data)
    return result


def unsorted_segment_mean(data, segment_ids, num_segments):
    '''
    :param data: [n_edge, *dimensions]
    :param segment_ids: [n_edge]
    :param num_segments: [bs * n_node]
    '''
    expand_dims = tuple(data.shape[1:])
    result_shape = (num_segments, ) + expand_dims
    for _ in expand_dims:
        segment_ids = segment_ids.unsqueeze(-1)
    segment_ids = segment_ids.expand(-1, *expand_dims)
    result = data.new_full(result_shape, 0)  # Init empty result tensor.
    count = data.new_full(result_shape, 0)
    result.scatter_add_(0, segment_ids, data)
    count.scatter_add_(0, segment_ids, torch.ones_like(data))
    return result / count.clamp(min=1)


def get_edges(n_nodes):
    rows, cols = [], []
    for i in range(n_nodes):
        for j in range(n_nodes):
            if i != j:
                rows.append(i)
                cols.append(j)

    edges = [rows, cols]
    return edges


def get_edges_batch(n_nodes, batch_size):
    edges = get_edges(n_nodes)
    edge_attr = torch.ones(len(edges[0]) * batch_size, 1)
    edges = [torch.LongTensor(edges[0]), torch.LongTensor(edges[1])]
    if batch_size == 1:
        return edges, edge_attr
    elif batch_size > 1:
        rows, cols = [], []
        for i in range(batch_size):
            rows.append(edges[0] + n_nodes * i)
            cols.append(edges[1] + n_nodes * i)
        edges = [torch.cat(rows), torch.cat(cols)]
    return edges, edge_attr


if __name__ == "__main__":
    # Dummy parameters
    batch_size = 8
    n_nodes = 4
    n_feat = 10
    x_dim = 3
    n_channel = 5

    # Dummy variables h, x and fully connected edges
    h = torch.randn(batch_size *  n_nodes, n_feat)
    x = torch.randn(batch_size * n_nodes, n_channel, x_dim)
    edges, edge_attr = get_edges_batch(n_nodes, batch_size)
    ctx_edges, att_edges = edges, edges

    # Initialize EGNN
    gnn = MCAttEGNN(in_node_nf=n_feat, hidden_nf=32, out_node_nf=21, n_channel=n_channel)

    # Run EGNN
    h, x = gnn(h, x, ctx_edges, att_edges)

    print(h)
    print(x)
AbDesign/diffab/modules/MCAttGNN/mc_egnn.py
pengzhangzhi/ab_opt
#!/usr/bin/python
# -*- coding:utf-8 -*-
from collections import defaultdict
from copy import copy, deepcopy
import json
import math
import os
from typing import Dict, List, Optional, Tuple, Union
import requests

import numpy as np
from Bio.PDB import PDBParser, PDBIO
from Bio.PDB.Structure import Structure as BStructure
from Bio.PDB.Model import Model as BModel
from Bio.PDB.Chain import Chain as BChain
from Bio.PDB.Residue import Residue as BResidue
from Bio.PDB.Atom import Atom as BAtom
from tqdm import tqdm



def rotation_matrix_from_vectors(vec1, vec2):
    """ Find the rotation matrix that aligns vec1 to vec2
    :param vec1: A 3d "source" vector
    :param vec2: A 3d "destination" vector
    :return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2.
    """
    a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3)
    v = np.cross(a, b)
    c = np.dot(a, b)
    s = np.linalg.norm(v)
    kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
    rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2))
    return rotation_matrix  # mat.dot(vec1) == vec2


# Gram-Schmidt Orthogonalization
def gram_schmidt(vectors):
    basis = []
    for v in vectors:
        w = v - np.sum([np.dot(v, b) * b  for b in basis], axis=0)
        if (np.abs(w) > 1e-10).any():  
            basis.append(w / np.linalg.norm(w))
    return np.array(basis)


class AminoAcid:
    def __init__(self, symbol, abrv, idx=0, side_chain_coord=None):
        self.symbol = symbol
        self.abrv = abrv
        self.idx = idx
        self.side_chain_coord = side_chain_coord

    def set_side_chain_coord(self, side_chain_coord):
        '''side_chain_coord:
        relative positions in the local coordinate with (CB-CA, N-CA, C-CA) (after gram schmidt)
        '''
        self.side_chain_coord = deepcopy(side_chain_coord)
        for atom in self.side_chain_coord:
            self.side_chain_coord[atom] = np.array(self.side_chain_coord[atom])

    def __str__(self):
        return f'{self.idx} {self.symbol} {self.abrv}'


class AminoAcidVocab:
    def __init__(self):
        self.PAD, self.SEP, self.UNK = '#', '/', '*'
        self.BOA, self.BOH, self.BOL = '&', '+', '-' # begin of antigen, heavy chain, light chain
        specials = [# special added (PAD)
                (self.PAD, 'PAD'), (self.UNK, 'UNK'),  # for RefineGNN, pad is mask
                (self.BOA, '<X>'), (self.BOH, '<H>'), (self.BOL, '<L>'), (self.SEP, '<E>')
            ]
        aas = [
                ('G', 'GLY'), ('A', 'ALA'), ('V', 'VAL'), ('L', 'LEU'),
                ('I', 'ILE'), ('F', 'PHE'), ('W', 'TRP'), ('Y', 'TYR'),
                ('D', 'ASP'), ('H', 'HIS'), ('N', 'ASN'), ('E', 'GLU'),
                ('K', 'LYS'), ('Q', 'GLN'), ('M', 'MET'), ('R', 'ARG'),
                ('S', 'SER'), ('T', 'THR'), ('C', 'CYS'), ('P', 'PRO'),
                ('U', 'SEC') # 21 aa for eukaryote
            ]
        _all = specials + aas
        self.amino_acids = [AminoAcid(symbol, abrv) for symbol, abrv in _all]
        self.symbol2idx, self.abrv2idx = {}, {}
        for i, aa in enumerate(self.amino_acids):
            self.symbol2idx[aa.symbol] = i
            self.abrv2idx[aa.abrv] = i
            aa.idx = i
        self.special_mask = [1 for _ in specials] + [0 for _ in aas]
        self.side_chain_loaded = False
    
    def load_side_chain_coord(self, path):
        self.side_chain_loaded = True
        return
        # not implemented
        with open(path, 'r') as fin:
            coord = json.load(fin)
        for symbol in coord:
            self.amino_acids[self.symbol_to_idx(symbol)].set_side_chain_coord(coord[symbol])
            # print(f'{symbol} side chain information loaded')
        print(f'{len(coord)} side chain information loaded')

    def abrv_to_symbol(self, abrv):
        idx = self.abrv_to_idx(abrv)
        return None if idx is None else self.amino_acids[idx].symbol

    def symbol_to_abrv(self, symbol):
        idx = self.symbol_to_idx(symbol)
        return None if idx is None else self.amino_acids[idx].abrv

    def abrv_to_idx(self, abrv):
        abrv = abrv.upper()
        return self.abrv2idx.get(abrv, None)

    def symbol_to_idx(self, symbol):
        symbol = symbol.upper()
        return self.symbol2idx.get(symbol, None)
    
    def idx_to_symbol(self, idx):
        return self.amino_acids[idx].symbol

    def idx_to_abrv(self, idx):
        return self.amino_acids[idx].abrv

    def get_pad_idx(self):
        return self.symbol_to_idx(self.PAD)

    def get_unk_idx(self):
        return self.symbol_to_idx(self.UNK)
    
    def get_special_mask(self):
        return copy(self.special_mask)

    def get_side_chain_info(self, symbol):
        assert self.side_chain_loaded, 'Side chain information not loaded!'
        symbol = self.symbol_to_idx(symbol)
        coord = self.amino_acids[symbol].side_chain_coord
        return deepcopy(coord)

    def __len__(self):
        return len(self.symbol2idx)


VOCAB = AminoAcidVocab()
# SIDE_CHAIN_FILE = os.path.join(os.path.split(__file__)[0], 'side_chain_coords.json')
# if os.path.exists(SIDE_CHAIN_FILE):
#     VOCAB.load_side_chain_coord(SIDE_CHAIN_FILE)
# else:
#     print(f'Side chain coordination file not found at {SIDE_CHAIN_FILE}')


def format_aa_abrv(abrv):  # special cases
    if abrv == 'MSE':
        return 'MET' # substitue MSE with MET
    return abrv


class Residue:
    def __init__(self, symbol: str, coordinate: Tuple, _id: Tuple):
        self.symbol = symbol
        self.coordinate = coordinate
        self.id = _id  # (residue_number, insert_code)

    def get_symbol(self):
        return self.symbol

    def get_coord(self, atom_name):
        return copy(self.coordinate[atom_name])

    def get_coord_map(self) -> Dict[str, List]:
        return deepcopy(self.coordinate)

    def get_backbone_coord_map(self) -> Dict[str, List]:
        coord = { atom: self.coordinate[atom] for atom in self.coordinate if atom in ['CA', 'C', 'N', 'O'] }
        return coord

    def get_side_chain_coord_map(self) -> Dict[str, List]:
        coord = { atom: self.coordinate[atom] for atom in self.coordinate if atom not in ['CA', 'C', 'N', 'O'] }
        return coord

    def get_atom_names(self):
        return list(self.coordinate.keys())

    def get_id(self):
        return self.id

    def set_symbol(self, symbol):
        assert VOCAB.symbol_to_abrv(symbol) is not None, f'{symbol} is not an amino acid'
        self.symbol = symbol

    def set_coord(self, coord):
        self.coordinate = deepcopy(coord)

    def gen_side_chain(self, center=None):
        '''
        automatically generate side chain according to statistics
        '''
        coord = VOCAB.get_side_chain_info(self.symbol)
        if coord is None:  # do not have a side chain
            return
        n, ca = np.array(self.get_coord('N')), np.array(self.get_coord('CA'))
        if center is None:
            center = coord['center'] + ca
        else:
            center = np.array(center)
        coord.pop('center')

        center_ca, n_ca = center - ca, n - ca
        X = gram_schmidt([center_ca, n_ca, np.cross(center_ca, n_ca)])
        for atom in coord:
            self.coordinate[atom] = (coord[atom].dot(X) + ca).tolist()

    def dist_to(self, residue):  # measured by nearest atoms
        xa = np.array(list(self.get_coord_map().values()))
        xb = np.array(list(residue.get_coord_map().values()))
        if len(xa) == 0 or len(xb) == 0:
            return math.nan
        dist = np.linalg.norm(xa[:, None, :] - xb[None, :, :], axis=-1)
        return np.min(dist)

    def to_bio(self):
        _id = (' ', self.id[0], self.id[1])
        residue = BResidue(_id, VOCAB.symbol_to_abrv(self.symbol), '    ')
        atom_map = self.coordinate
        for i, atom in enumerate(atom_map):
            fullname = ' ' + atom
            while len(fullname) < 4:
                fullname += ' '
            bio_atom = BAtom(
                name=atom,
                coord=np.array(atom_map[atom], dtype=np.float32),
                bfactor=0,
                occupancy=1.0,
                altloc=' ',
                fullname=fullname,
                serial_number=i,
                element=atom[0]  # not considering symbols with 2 chars (e.g. FE, MG)
            )
            residue.add(bio_atom)
        return residue


class Peptide:
    def __init__(self, _id, residues: List[Residue]):
        self.residues = residues
        self.seq = ''
        self.id = _id
        for residue in residues:
            self.seq += residue.get_symbol()

    def set_id(self, _id):
        self.id = _id

    def get_id(self):
        return self.id

    def get_seq(self):
        return self.seq

    def get_span(self, i, j):  # [i, j)
        i, j = max(i, 0), min(j, len(self.seq))
        if j <= i:
            residues = []
        else:
            residues = deepcopy(self.residues[i:j])
        return Peptide(self.id, residues)

    def get_residue(self, i):
        return deepcopy(self.residues[i])
    
    def get_ca_pos(self, i):
        return copy(self.residues[i].get_coord('CA'))

    def get_cb_pos(self, i):
        return copy(self.residues[i].get_coord('CB'))

    def set_residue_coord(self, i, coord):
        self.residues[i].set_coord(coord)

    def set_residue_translation(self, i, vec):
        coord = self.residues[i].get_coord_map()
        for atom in coord:
            ori_vec = coord[atom]
            coord[atom] = [a + b for a, b in zip(ori_vec, vec)]
        self.set_residue_coord(i, coord)

    def set_residue_symbol(self, i, symbol):
        self.residues[i].set_symbol(symbol)
        self.seq = self.seq[:i] + symbol + self.seq[i+1:]

    def set_residue(self, i, symbol, coord, center=None, gen_side_chain=False):
        self.set_residue_symbol(i, symbol)
        self.set_residue_coord(i, coord)
        if gen_side_chain:
            self.residues[i].gen_side_chain(center)

    def to_bio(self):
        chain = BChain(id=self.id)
        for residue in self.residues:
            chain.add(residue.to_bio())
        return chain

    def __len__(self):
        return len(self.seq)

    def __str__(self):
        return self.seq


class Protein:
    def __init__(self, pdb_id, peptides):
        self.pdb_id = pdb_id
        self.peptides = peptides

    @classmethod
    def from_pdb(cls, pdb_path):
        parser = PDBParser(QUIET=True)
        structure = parser.get_structure('anonym', pdb_path)
        pdb_id = structure.header['idcode'].upper()
        peptides = {}
        for chain in structure.get_chains():
            _id = chain.get_id()
            residues = []
            has_non_residue = False
            for residue in chain:
                abrv = residue.get_resname()
                hetero_flag, res_number, insert_code = residue.get_id()
                if hetero_flag != ' ':
                    continue   # residue from glucose or water
                symbol = VOCAB.abrv_to_symbol(abrv)
                if symbol is None:
                    has_non_residue = True
                    print(f'has non residue: {abrv}')
                    break
                # filter Hs because not all data include them
                atoms = { atom.get_id(): atom.get_coord() for atom in residue if atom.element != 'H' }
                residues.append(Residue(
                    symbol, atoms, (res_number, insert_code)
                ))
            if has_non_residue or len(residues) == 0:  # not a peptide
                continue
            peptides[_id] = Peptide(_id, residues)
        return cls(pdb_id, peptides)

    def get_id(self):
        return self.pdb_id

    def num_chains(self):
        return len(self.peptides)

    def get_chain(self, name):
        if name in self.peptides:
            return deepcopy(self.peptides[name])
        else:
            return Peptide(name, [])

    def get_chain_names(self):
        return list(self.peptides.keys())

    def to_bio(self):
        structure = BStructure(id=self.pdb_id)
        model = BModel(id=0)
        for name in self.peptides:
            model.add(self.peptides[name].to_bio())
        structure.add(model)
        return structure

    def to_pdb(self, path):
        bio_structure = self.to_bio()
        io = PDBIO()
        io.set_structure(bio_structure)
        io.save(path)

    def __eq__(self, other):
        if not isinstance(other, Protein):
            raise TypeError('Cannot compare other type to Protein')
        for key in self.peptides:
            if key in other.peptides and self.peptides[key].seq == other.peptides[key].seq:
                continue
            else:
                return False
        return True

    def __str__(self):
        res = self.pdb_id + '\n'
        for seq_name in self.peptides:
            res += f'\t{seq_name}: {self.peptides[seq_name]}\n'
        return res


def continuous_segments(indexes):
    segs, indexes = [], sorted(list(set(indexes)))
    return indexes
    print(len(indexes))
    cur_seg = None
    for i in indexes:
        if cur_seg is None:
            cur_seg = [i, i]
        elif i <= cur_seg[1]:
            cur_seg[1] = i
        else:  # finish a segment
            if cur_seg[0] != cur_seg[1]:  # drop single residue
                segs.append(cur_seg)
            cur_seg = [i, i]
    if cur_seg is not None and cur_seg[0] != cur_seg[1]:
        segs.append(cur_seg)
    return segs


class AAComplex(Protein):  # Antibody-Antigen complex

    # threshold = 6.6  # from "Characterization of Protein-Protein Interfaces"
    num_interface_residue = 48  # from PNAS, Jian Peng

    def __init__(self, pdb_id: str, peptides: Dict[str, Peptide], heavy_chain: str,
                 light_chain: str, antigen_chains: List[str], numbering: str='imgt',
                 cdr_pos: Optional[Dict[str, Tuple]]=None, skip_cal_interface=False):
        '''
            heavy_chain: the id of heavy chain
            light_chain: the id of light chain
            antigen_chains: the list of ids of antigen chains
            numbering: currently only support IMGT
            self.cdr_pos: dict with keys like CDR-H3 and values like (40, 45) indicating [begin, end]
        '''
        self.heavy_chain = heavy_chain
        self.light_chain = light_chain
        self.antigen_chains = copy(antigen_chains)
    
        # antibody information
        if cdr_pos is None:
            selected_peptides, self.cdr_pos = self._extract_antibody_info(peptides, numbering)
        else:
            selected_peptides, self.cdr_pos = {}, deepcopy(cdr_pos)
            for chain_name in [heavy_chain, light_chain]:
                if chain_name in peptides:
                    selected_peptides[chain_name] = peptides[chain_name]
        # antigen information
        for chain_name in antigen_chains:
            assert chain_name not in selected_peptides, f'Antigen chain {chain_name} is antibody itself!'
            selected_peptides[chain_name] = peptides[chain_name]

        super().__init__(pdb_id, selected_peptides)

        if not skip_cal_interface:
            self._cal_interface()

    def _extract_antibody_info(self, peptides, numbering):
        # calculating cdr pos according to number scheme (type_mapping and conserved residues)
        if numbering.lower() == 'imgt':
            type_mapping = {}  # - for non-Fv region, 0 for framework, 1/2/3 for cdr1/2/3
            for i in list(range(1, 27)) + list(range(39, 56)) + list(range(66, 105)) + list(range(118, 130)):
                type_mapping[i] = '0'
            for i in range(27, 39):     # cdr1
                type_mapping[i] = '1'
            for i in range(56, 66):     # cdr2
                type_mapping[i] = '2'
            for i in range(105, 118):   # cdr3
                type_mapping[i] = '3'
            conserved = {
                23: ['CYS'],
                41: ['TRP'],
                104: ['CYS'],
                # 118: ['PHE', 'TRP']
            }
        else:
            raise NotImplementedError(f'Numbering scheme {numbering} not implemented')

        selected_peptides, cdr_pos = {}, {}
        for c, chain_name in zip(['H', 'L'], [self.heavy_chain, self.light_chain]):
            chain = None
            for _name in [chain_name, chain_name.upper()]:
                # Note: possbly two chains are different segments of a same chain
                if _name in peptides:
                    chain = peptides[_name]
                    break
            if chain is None:
                continue           
            res_type = ''
            for i in range(len(chain)):
                residue = chain.get_residue(i)
                residue_number = residue.get_id()[0]
                if residue_number in type_mapping:
                    res_type += type_mapping[residue_number]
                    if residue_number in conserved:
                        hit, symbol = False, residue.get_symbol()
                        for conserved_residue in conserved[residue_number]:
                            if symbol == VOCAB.abrv_to_symbol(conserved_residue):
                                hit = True
                                break
                        assert hit, f'Not {conserved[residue_number]} at {residue_number}'
                else:
                    res_type += '-'
            start, end = res_type.index('0'), res_type.rindex('0')
            for cdr in ['1', '2', '3']:
                cdr_start, cdr_end = res_type.find(cdr), res_type.rfind(cdr)
                if cdr_start == -1:
                    raise ValueError(f'cdr {cdr} not found, residue type: {res_type}')
                start, end = min(start, cdr_start), max(end, cdr_end)
                cdr_pos[f'CDR-{c}{cdr}'] = (cdr_start, cdr_end)
            for cdr in ['1', '2', '3']:
                cdr = f'CDR-{c}{cdr}'
                cdr_start, cdr_end = cdr_pos[cdr]
                cdr_pos[cdr] = (cdr_start - start, cdr_end - start)
            chain = chain.get_span(start, end + 1)  # the length may exceed 130 because of inserted amino acids
            chain.set_id(chain_name)
            selected_peptides[chain_name] = chain

        return selected_peptides, cdr_pos

    def _cal_interface(self):
        antigen_chains = self.get_antigen_chains()
        antigen_names = [peptide.get_id() for peptide in antigen_chains]
        antibody_chains = [self.get_heavy_chain(), self.get_light_chain()]
        antibody_chains = [chain for chain in antibody_chains if chain is not None]
        antibody_names = [chain.get_id() for chain in antibody_chains]
        coord = {}
        for name, chain in zip(antigen_names + antibody_names, antigen_chains + antibody_chains):
            coord[name] = {'x': [], 'i': [], 'interface': []}
            for i in range(len(chain)):
                try:
                    coord[name]['x'].append(chain.get_ca_pos(i))
                except KeyError:
                    print(f'{self.pdb_id}, chain {chain.get_id()}, residue at {chain.get_residue(i).get_id()} has no ca coordination, fill with other atoms')
                    coord_map = chain.get_residue(i).get_coord_map()
                    atom = list(coord_map.keys())[0]
                    coord[name]['x'].append(coord_map[atom])
                # coord[name]['i'].append(i)
                # residue = chain.get_residue(i)
                # for atom in residue.get_atom_names():
                #     x = residue.get_coord(atom)
                #     coord[name]['x'].append(x)
                #     coord[name]['i'].append(i)

        for name in coord:
            coord[name]['x'] = np.array(coord[name]['x'])
            # coord[name]['i'] = np.array(coord[name]['i'])

        # calculate distance
        min_dists = { name: [] for name in antigen_names }
        for name in antigen_names:
            xa = coord[name]['x']
            for ab_name in antibody_names:
                xb = coord[ab_name]['x']
                if len(xb) == 0:  # for single-domain antibodies
                    continue
                dist = np.linalg.norm(xa[:, None, :] - xb[None, :, :], axis=-1)
                min_dist = np.min(dist, axis=1) # [num_ag_residue]
                min_dists[name].append(min_dist)
                # _is, _js = np.nonzero(dist <= self.threshold)
                # _is, _js = coord[name]['i'][_is], coord[ab_name]['i'][_js]
                # for n, indexes in zip([name, ab_name], [_is, _js]):
                #     # if indexes.shape[0] == 0:
                #     #     break
                #     # print(indexes)
                #     # start, end = np.min(indexes), np.max(indexes)
                #     # first put all index into it
                #     coord[n]['interface'].extend(set(indexes))
                #     # coord[n]['interface'][0] = min(coord[n]['interface'][0], start)
                #     # coord[n]['interface'][1] = max(coord[n]['interface'][1], end)
        dists, ids = [], []
        for name in antigen_names:
            if len(min_dists[name]) == 2:
                min_dists[name] = np.minimum(*min_dists[name])
            elif len(min_dists[name]) == 1:
                min_dists[name] = min_dists[name][0]
            else:
                raise ValueError('Number of chain in ab exceed 2')
            dists.extend(min_dists[name])
            for i in range(len(min_dists[name])):
                ids.append((name, i))
        dists = np.array(dists)
        topk = min(len(dists), self.num_interface_residue)
        ind = np.argpartition(-dists, -topk)[-topk:]
        self.antigen_interface = {}
        for i in ind:
            name, residue_idx = ids[i]
            if name not in self.antigen_interface:
                self.antigen_interface[name] = []
            self.antigen_interface[name].append(residue_idx)

        # self.antigen_interface = { name: continuous_segments(coord[name]['interface']) for name in antigen_names }
        # self.antibody_interface = { name: continuous_segments(coord[name]['interface']) for name in antibody_names }
        # self.antigen_cdr_interface = {}

        # for symbol, chain_name in zip(['H', 'L'], [self.heavy_chain, self.light_chain]):
        #     for i in range(1, 4):
        #         start, end = self.get_cdr_pos(cdr=symbol + str(i))
        #         if start >= end:
        #             continue
        #         for name in antigen_names:
        #             xa = coord[name]['x']
        #             xb = coord[chain_name]['x'][start:end + 1]
        #             dist = np.linalg.norm(xa[:, None, :] - xb[None, :, :], axis=-1)
        #             _is, _ = np.nonzero(dist <= self.threshold)
        #             if len(_is) == 0:
        #                 continue
        #             _is = coord[name]['i'][_is]
        #             # start, end = np.min(_is), np.max(_is)
        #             # self.antigen_cdr_interface[f'{name}-CDR-{symbol}{i}'] = [start, end]
        #             self.antigen_cdr_interface[f'{name}-CDR-{symbol}{i}'] = continuous_segments(_is)

    def get_heavy_chain(self, interface_only=False) -> Union[Peptide, List[Peptide]]:
        chain = self.get_chain(self.heavy_chain)
        if not len(chain):
            return chain
        if interface_only:
            raise NotImplementedError('get heavy chain interface not implemented')
            spans = self.antibody_interface[self.heavy_chain]
            chain = Peptide(chain.get_id(), [chain.get_residue(i) for i in spans])
        return chain

    def get_light_chain(self, interface_only=False) -> Union[Peptide, List[Peptide]]:
        chain = self.get_chain(self.light_chain)
        if not len(chain):
            return chain
        if interface_only:
            raise NotImplementedError('get light chain interface not implemented')
            spans = self.antibody_interface[self.light_chain]
            chain = [chain.get_span(start, end + 1) for start, end in spans]
        return chain

    def get_antigen_chains(self, interface_only=False, cdr=None) -> List[Peptide]: # H/L + 1/2/3, None for the whole antibody
        chains = [self.get_chain(name) for name in self.antigen_chains]
        if interface_only:
            new_chains = []
            cdr = None if cdr is None else f'CDR-{cdr.upper()}'
            for chain in chains:
                if cdr is None:
                    spans = self.antigen_interface.get(chain.get_id(), [])
                else:
                    name = f'{chain.get_id()}-{cdr}'
                    spans = self.antigen_cdr_interface.get(name, [])
                # for span in spans:
                #     new_chains.append(chain.get_span(span[0], span[1] + 1))
                new_chains.append(Peptide(chain.get_id(), [chain.get_residue(i) for i in spans]))
            chains = [chain for chain in new_chains if len(chain)]
        return chains

    def get_cdr_pos(self, cdr='H3'):  # H/L + 1/2/3, return [begin, end] position
        cdr = f'CDR-{cdr}'.upper()
        if cdr in self.cdr_pos:
            return self.cdr_pos[cdr]
        else:
            return (-1, -1)

    def get_cdr(self, cdr='H3'):
        cdr = cdr.upper()
        pos = self.get_cdr_pos(cdr)
        chain = self.get_heavy_chain() if 'H' in cdr else self.get_light_chain()
        return chain.get_span(pos[0], pos[1] + 1)

    def __str__(self):
        pdb_info = f'PDB ID: {self.pdb_id}'
        antibody_info = f'Antibody H-{self.heavy_chain} ({len(self.get_heavy_chain())}), ' + \
                        f'L-{self.light_chain} ({len(self.get_light_chain())})'
        antigen_info = f'Antigen Chains: {[(ag, len(self.get_chain(ag))) for ag in self.antigen_chains]}'
        cdr_info = f'CDRs: \n'
        for name in self.cdr_pos:
            chain = self.get_heavy_chain() if 'H' in name else self.get_light_chain()
            start, end = self.cdr_pos[name]
            cdr_info += f'\t{name}: [{start}, {end}], {chain.seq[start:end + 1]}\n'
        sep = '\n' + '=' * 20 + '\n'
        return sep + pdb_info + '\n' + antibody_info + '\n' + cdr_info + '\n' + antigen_info + sep


def construct_side_chain_coord(pdb_dir, out_path):
    side_chain_coord = {}
    # line = 0
    for f in tqdm(os.listdir(pdb_dir)):
        # line += 1
        # if line > 20:
        #     break
        try:
            protein = Protein.from_pdb(os.path.join(pdb_dir, f))
        except Exception as e:
            print(f'{e}, skip')
            continue
        for c in protein.get_chain_names():
            chain = protein.get_chain(c)
            for i in range(len(chain)):
                residue = chain.get_residue(i)
                atom_coord = residue.get_coord_map()
                try:
                    ca, n = atom_coord['CA'], atom_coord['N']
                    c = atom_coord['C']
                except KeyError:  # coordinates are missing
                    continue
                side_chain_atom_coord = residue.get_side_chain_coord_map()
                if len(side_chain_atom_coord) < 2:  # only one CB or no heavy atoms (GLY)
                    continue
                center = [side_chain_atom_coord[atom] for atom in side_chain_atom_coord]
                center = np.mean(center, axis=0)
                center_ca, n_ca = np.array([center, n]) - np.array([ca, ca])
                X, ca = gram_schmidt([center_ca, n_ca, np.cross(center_ca, n_ca)]), np.array(ca)
                symbol = residue.get_symbol()
                if symbol not in side_chain_coord:
                    side_chain_coord[symbol] = defaultdict(list)
                side_chain_coord[symbol]['center'].append(center_ca)
                for atom in side_chain_atom_coord:
                    x = np.array(atom_coord[atom]) - ca
                    assert len(x) == 3, x
                    x = X.dot(x)  # coordinate in the local frame
                    assert len(x) == 3, f'local: {x}, {X}, {center_ca}, {n_ca}'
                    side_chain_coord[symbol][atom].append(x)
    for symbol in side_chain_coord:
        print(f'{symbol}:')
        atom_coords = {}
        for atom in side_chain_coord[symbol]:
            coord = np.mean(side_chain_coord[symbol][atom], axis=0).tolist()
            std = np.std(side_chain_coord[symbol][atom], axis=0)
            print(f'\t{atom}: {coord}, {std}')
            atom_coords[atom] = coord
        side_chain_coord[symbol] = atom_coords
    with open(out_path, 'w') as fout:
        json.dump(side_chain_coord, fout)


def fetch_from_pdb(identifier):
    # example identifier: 1FBI
    url = 'https://data.rcsb.org/rest/v1/core/entry/' + identifier
    res = requests.get(url)
    if res.status_code != 200:
        return None
    url = f'https://files.rcsb.org/download/{identifier}.pdb'
    text = requests.get(url)
    data = res.json()
    data['pdb'] = text.text
    return data


if __name__ == '__main__':
    import sys
    construct_side_chain_coord(sys.argv[1], sys.argv[2])
AbDesign/diffab/modules/MCAttGNN/pdb_utils.py
pengzhangzhi/ab_opt
import torch
import torch.nn.functional as F

from diffab.utils.protein.constants import (
    BBHeavyAtom, 
    backbone_atom_coordinates_tensor,
    bb_oxygen_coordinate_tensor,
)
from .topology import get_terminus_flag


def safe_norm(x, dim=-1, keepdim=False, eps=1e-8, sqrt=True):
    out = torch.clamp(torch.sum(torch.square(x), dim=dim, keepdim=keepdim), min=eps)
    return torch.sqrt(out) if sqrt else out


def pairwise_distances(x, y=None, return_v=False):
    """
    Args:
        x:  (B, N, d)
        y:  (B, M, d)
    """
    if y is None: y = x
    v = x.unsqueeze(2) - y.unsqueeze(1)  # (B, N, M, d)
    d = safe_norm(v, dim=-1)
    if return_v:
        return d, v
    else:
        return d


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


def project_v2v(v, e, dim):
    """
    Description:
        Project vector `v` onto vector `e`.
    Args:
        v:  (N, L, 3).
        e:  (N, L, 3).
    """
    return (e * v).sum(dim=dim, keepdim=True) * e


def construct_3d_basis(center, p1, p2):
    """
    Args:
        center: (N, L, 3), usually the position of C_alpha.
        p1:     (N, L, 3), usually the position of C.
        p2:     (N, L, 3), usually the position of N.
    Returns
        A batch of orthogonal basis matrix, (N, L, 3, 3cols_index).
        The matrix is composed of 3 column vectors: [e1, e2, e3].
    """
    v1 = p1 - center    # (N, L, 3)
    e1 = normalize_vector(v1, dim=-1)

    v2 = p2 - center    # (N, L, 3)
    u2 = v2 - project_v2v(v2, e1, dim=-1)
    e2 = normalize_vector(u2, dim=-1)

    e3 = torch.cross(e1, e2, dim=-1)    # (N, L, 3)

    mat = torch.cat([
        e1.unsqueeze(-1), e2.unsqueeze(-1), e3.unsqueeze(-1)
    ], dim=-1)  # (N, L, 3, 3_index)
    return mat


def local_to_global(R, t, p):
    """
    Description:
        Convert local (internal) coordinates to global (external) coordinates q.
        q <- Rp + t
    Args:
        R:  (N, L, 3, 3).
        t:  (N, L, 3).
        p:  Local coordinates, (N, L, ..., 3).
    Returns:
        q:  Global coordinates, (N, L, ..., 3).
    """
    assert p.size(-1) == 3
    p_size = p.size()
    # N, L = p_size[0], p_size[1]

    p = p.transpose(-1, -2)   # (N, L, *, 3) -> (N, L, 3, *)
    q = torch.matmul(R, p) + t.unsqueeze(-1)    # (N, L, 3, *)
    q = q.transpose(-1, -2).reshape(p_size)     # (N, L, 3, *) -> (N, L, *, 3) -> (N, L, ..., 3)
    return q


def global_to_local(R, t, q):
    """
    Description:
        Convert global (external) coordinates q to local (internal) coordinates p.
        p <- R^{T}(q - t)
    Args:
        R:  (N, L, 3, 3).
        t:  (N, L, 3).
        q:  Global coordinates, (N, L, ..., 3).
    Returns:
        p:  Local coordinates, (N, L, ..., 3).
    """
    assert q.size(-1) == 3
    q_size = q.size()
    N, L = q_size[0], q_size[1]

    q = q.reshape(N, L, -1, 3).transpose(-1, -2)   # (N, L, *, 3) -> (N, L, 3, *)
    p = torch.matmul(R.transpose(-1, -2), (q - t.unsqueeze(-1)))  # (N, L, 3, *)
    p = p.transpose(-1, -2).reshape(q_size)     # (N, L, 3, *) -> (N, L, *, 3) -> (N, L, ..., 3)
    return p


def apply_rotation_to_vector(R, p):
    return local_to_global(R, torch.zeros_like(p), p)


def compose_rotation_and_translation(R1, t1, R2, t2):
    """
    Args:
        R1,t1:  Frame basis and coordinate, (N, L, 3, 3), (N, L, 3).
        R2,t2:  Rotation and translation to be applied to (R1, t1), (N, L, 3, 3), (N, L, 3).
    Returns
        R_new <- R1R2
        t_new <- R1t2 + t1
    """
    R_new = torch.matmul(R1, R2)    # (N, L, 3, 3)
    t_new = torch.matmul(R1, t2.unsqueeze(-1)).squeeze(-1) + t1
    return R_new, t_new


def compose_chain(Ts):
    while len(Ts) >= 2:
        R1, t1 = Ts[-2]
        R2, t2 = Ts[-1]
        T_next = compose_rotation_and_translation(R1, t1, R2, t2)
        Ts = Ts[:-2] + [T_next]
    return Ts[0]


# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
def quaternion_to_rotation_matrix(quaternions):
    """
    Convert rotations given as quaternions to rotation matrices.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    quaternions = F.normalize(quaternions, dim=-1)
    r, i, j, k = torch.unbind(quaternions, -1)
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    return o.reshape(quaternions.shape[:-1] + (3, 3))


# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""
BSD License

For PyTorch3D software

Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

 * Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.

 * Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

 * Neither the name Meta nor the names of its contributors may be used to
   endorse or promote products derived from this software without specific
   prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
def quaternion_1ijk_to_rotation_matrix(q):
    """
    (1 + ai + bj + ck) -> R
    Args:
        q:  (..., 3)
    """
    b, c, d = torch.unbind(q, dim=-1)
    s = torch.sqrt(1 + b**2 + c**2 + d**2)
    a, b, c, d = 1/s, b/s, c/s, d/s

    o = torch.stack(
        (
            a**2 + b**2 - c**2 - d**2,  2*b*c - 2*a*d,  2*b*d + 2*a*c,
            2*b*c + 2*a*d,  a**2 - b**2 + c**2 - d**2,  2*c*d - 2*a*b,
            2*b*d - 2*a*c,  2*c*d + 2*a*b,  a**2 - b**2 - c**2 + d**2,
        ),
        -1,
    )
    return o.reshape(q.shape[:-1] + (3, 3))


def repr_6d_to_rotation_matrix(x):
    """
    Args:
        x:  6D representations, (..., 6).
    Returns:
        Rotation matrices, (..., 3, 3_index).
    """
    a1, a2 = x[..., 0:3], x[..., 3:6]
    b1 = normalize_vector(a1, dim=-1)
    b2 = normalize_vector(a2 - project_v2v(a2, b1, dim=-1), dim=-1)
    b3 = torch.cross(b1, b2, dim=-1)

    mat = torch.cat([
        b1.unsqueeze(-1), b2.unsqueeze(-1), b3.unsqueeze(-1)
    ], dim=-1)  # (N, L, 3, 3_index)
    return mat


def dihedral_from_four_points(p0, p1, p2, p3):
    """
    Args:
        p0-3:   (*, 3).
    Returns:
        Dihedral angles in radian, (*, ).
    """
    v0 = p2 - p1
    v1 = p0 - p1
    v2 = p3 - p2
    u1 = torch.cross(v0, v1, dim=-1)
    n1 = u1 / torch.linalg.norm(u1, dim=-1, keepdim=True)
    u2 = torch.cross(v0, v2, dim=-1)
    n2 = u2 / torch.linalg.norm(u2, dim=-1, keepdim=True)
    sgn = torch.sign( (torch.cross(v1, v2, dim=-1) * v0).sum(-1) )
    dihed = sgn*torch.acos( (n1 * n2).sum(-1).clamp(min=-0.999999, max=0.999999) )
    dihed = torch.nan_to_num(dihed)
    return dihed


def knn_gather(idx, value):
    """
    Args:
        idx:    (B, N, K)
        value:  (B, M, d)
    Returns:
        (B, N, K, d)
    """
    N, d = idx.size(1), value.size(-1)
    idx = idx.unsqueeze(-1).repeat(1, 1, 1, d)      # (B, N, K, d)
    value = value.unsqueeze(1).repeat(1, N, 1, 1)   # (B, N, M, d)
    return torch.gather(value, dim=2, index=idx)


def knn_points(q, p, K):
    """
    Args:
        q: (B, M, d)
        p: (B, N, d)
    Returns:
        (B, M, K), (B, M, K), (B, M, K, d)
    """
    _, L, _ = p.size()
    d = pairwise_distances(q, p)  # (B, N, M)
    dist, idx = d.topk(min(L, K), dim=-1, largest=False)  # (B, M, K), (B, M, K)
    return dist, idx, knn_gather(idx, p)


def angstrom_to_nm(x):
    return x / 10


def nm_to_angstrom(x):
    return x * 10


def get_backbone_dihedral_angles(pos_atoms, chain_nb, res_nb, mask):
    """
    Args:
        pos_atoms:  (N, L, A, 3).
        chain_nb:   (N, L).
        res_nb:     (N, L).
        mask:       (N, L).
    Returns:
        bb_dihedral:    Omega, Phi, and Psi angles in radian, (N, L, 3).
        mask_bb_dihed:  Masks of dihedral angles, (N, L, 3).
    """
    pos_N  = pos_atoms[:, :, BBHeavyAtom.N]   # (N, L, 3)
    pos_CA = pos_atoms[:, :, BBHeavyAtom.CA]
    pos_C  = pos_atoms[:, :, BBHeavyAtom.C]

    N_term_flag, C_term_flag = get_terminus_flag(chain_nb, res_nb, mask)  # (N, L)
    omega_mask = torch.logical_not(N_term_flag)
    phi_mask = torch.logical_not(N_term_flag)
    psi_mask = torch.logical_not(C_term_flag)

    # N-termini don't have omega and phi
    omega = F.pad(
        dihedral_from_four_points(pos_CA[:, :-1], pos_C[:, :-1], pos_N[:, 1:], pos_CA[:, 1:]), 
        pad=(1, 0), value=0,
    )
    phi = F.pad(
        dihedral_from_four_points(pos_C[:, :-1], pos_N[:, 1:], pos_CA[:, 1:], pos_C[:, 1:]),
        pad=(1, 0), value=0,
    )

    # C-termini don't have psi
    psi = F.pad(
        dihedral_from_four_points(pos_N[:, :-1], pos_CA[:, :-1], pos_C[:, :-1], pos_N[:, 1:]),
        pad=(0, 1), value=0,
    )

    mask_bb_dihed = torch.stack([omega_mask, phi_mask, psi_mask], dim=-1)
    bb_dihedral = torch.stack([omega, phi, psi], dim=-1) * mask_bb_dihed
    return bb_dihedral, mask_bb_dihed


def pairwise_dihedrals(pos_atoms):
    """
    Args:
        pos_atoms:  (N, L, A, 3).
    Returns:
        Inter-residue Phi and Psi angles, (N, L, L, 2).
    """
    N, L = pos_atoms.shape[:2]
    pos_N  = pos_atoms[:, :, BBHeavyAtom.N]   # (N, L, 3)
    pos_CA = pos_atoms[:, :, BBHeavyAtom.CA]
    pos_C  = pos_atoms[:, :, BBHeavyAtom.C]

    ir_phi = dihedral_from_four_points(
        pos_C[:,:,None].expand(N, L, L, 3), 
        pos_N[:,None,:].expand(N, L, L, 3), 
        pos_CA[:,None,:].expand(N, L, L, 3), 
        pos_C[:,None,:].expand(N, L, L, 3)
    )
    ir_psi = dihedral_from_four_points(
        pos_N[:,:,None].expand(N, L, L, 3), 
        pos_CA[:,:,None].expand(N, L, L, 3), 
        pos_C[:,:,None].expand(N, L, L, 3), 
        pos_N[:,None,:].expand(N, L, L, 3)
    )
    ir_dihed = torch.stack([ir_phi, ir_psi], dim=-1)
    return ir_dihed


def apply_rotation_matrix_to_rot6d(R, O):
    """
    Args:
        R:  (..., 3, 3)
        O:  (..., 6)
    Returns:
        Rotated 6D representation, (..., 6).
    """
    u1, u2 = O[..., :3, None], O[..., 3:, None] # (..., 3, 1)
    v1 = torch.matmul(R, u1).squeeze(-1)    # (..., 3)
    v2 = torch.matmul(R, u2).squeeze(-1)
    return torch.cat([v1, v2], dim=-1)


def normalize_rot6d(O):
    """
    Args:
        O:  (..., 6)
    """
    u1, u2 = O[..., :3], O[..., 3:]     # (..., 3)
    v1 = F.normalize(u1, p=2, dim=-1)   # (..., 3)
    v2 = F.normalize(u2 - project_v2v(u2, v1), p=2, dim=-1)
    return torch.cat([v1, v2], dim=-1)
    

def reconstruct_backbone(R, t, aa, chain_nb, res_nb, mask):
    """
    Args:
        R:  (N, L, 3, 3)
        t:  (N, L, 3)
        aa: (N, L)
        chain_nb:   (N, L)
        res_nb:     (N, L)
        mask:       (N, L)
    Returns:
        Reconstructed backbone atoms, (N, L, 4, 3).
    """
    N, L = aa.size()
    # atom_coords = restype_heavyatom_rigid_group_positions.clone().to(t) # (21, 14, 3)
    bb_coords = backbone_atom_coordinates_tensor.clone().to(t)  # (21, 3, 3)
    oxygen_coord = bb_oxygen_coordinate_tensor.clone().to(t)    # (21, 3)
    aa = aa.clamp(min=0, max=20)    # 20 for UNK

    bb_coords = bb_coords[aa.flatten()].reshape(N, L, -1, 3)    # (N, L, 3, 3)
    oxygen_coord = oxygen_coord[aa.flatten()].reshape(N, L, -1)  # (N, L, 3)
    bb_pos = local_to_global(R, t, bb_coords)   # Global coordinates of N, CA, C. (N, L, 3, 3).

    # Compute PSI angle
    bb_dihedral, _ = get_backbone_dihedral_angles(bb_pos, chain_nb, res_nb, mask)
    psi = bb_dihedral[..., 2]   # (N, L)
    # Make rotation matrix for PSI
    sin_psi = torch.sin(psi).reshape(N, L, 1, 1)
    cos_psi = torch.cos(psi).reshape(N, L, 1, 1)
    zero = torch.zeros_like(sin_psi)
    one = torch.ones_like(sin_psi)
    row1 = torch.cat([one, zero, zero], dim=-1)     # (N, L, 1, 3)
    row2 = torch.cat([zero, cos_psi, -sin_psi], dim=-1) # (N, L, 1, 3)
    row3 = torch.cat([zero, sin_psi, cos_psi], dim=-1)  # (N, L, 1, 3)
    R_psi = torch.cat([row1, row2, row3], dim=-2)       # (N, L, 3, 3)

    # Compute rotoation and translation of PSI frame, and position of O.
    R_psi, t_psi = compose_chain([
        (R, t), # Backbone
        (R_psi, torch.zeros_like(t)),       # PSI angle
    ])
    O_pos = local_to_global(R_psi, t_psi, oxygen_coord.reshape(N, L, 1, 3))

    bb_pos = torch.cat([bb_pos, O_pos], dim=2)  # (N, L, 4, 3)
    return bb_pos
    

def reconstruct_backbone_partially(pos_ctx, R_new, t_new, aa, chain_nb, res_nb, mask_atoms, mask_recons):
    """
    Args:
        pos:    (N, L, A, 3).
        R_new:  (N, L, 3, 3).
        t_new:  (N, L, 3).
        mask_atoms: (N, L, A).
        mask_recons:(N, L).
    Returns:
        pos_new:    (N, L, A, 3).
        mask_new:   (N, L, A).
    """
    N, L, A = mask_atoms.size()

    mask_res = mask_atoms[:, :, BBHeavyAtom.CA]
    pos_recons = reconstruct_backbone(R_new, t_new, aa, chain_nb, res_nb, mask_res) # (N, L, 4, 3)
    pos_recons = F.pad(pos_recons, pad=(0, 0, 0, A-4), value=0) # (N, L, A, 3)

    pos_new = torch.where(
        mask_recons[:, :, None, None].expand_as(pos_ctx),
        pos_recons, pos_ctx
    )   # (N, L, A, 3)

    mask_bb_atoms = torch.zeros_like(mask_atoms)
    mask_bb_atoms[:, :, :4] = True
    mask_new = torch.where(
        mask_recons[:, :, None].expand_as(mask_atoms),
        mask_bb_atoms, mask_atoms
    )

    return pos_new, mask_new


AbDesign/diffab/modules/common/geometry.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F


def mask_zero(mask, value):
    return torch.where(mask, value, torch.zeros_like(value))


def clampped_one_hot(x, num_classes):
    mask = (x >= 0) & (x < num_classes) # (N, L)
    x = x.clamp(min=0, max=num_classes-1)
    y = F.one_hot(x, num_classes) * mask[...,None]  # (N, L, C)
    return y


class DistanceToBins(nn.Module):

    def __init__(self, dist_min=0.0, dist_max=20.0, num_bins=64, use_onehot=False):
        super().__init__()
        self.dist_min = dist_min
        self.dist_max = dist_max
        self.num_bins = num_bins
        self.use_onehot = use_onehot

        if use_onehot:
            offset = torch.linspace(dist_min, dist_max, self.num_bins)
        else:
            offset = torch.linspace(dist_min, dist_max, self.num_bins-1)    # 1 overflow flag
            self.coeff = -0.5 / ((offset[1] - offset[0]) * 0.2).item() ** 2  # `*0.2`: makes it not too blurred
        self.register_buffer('offset', offset)

    @property
    def out_channels(self):
        return self.num_bins 

    def forward(self, dist, dim, normalize=True):
        """
        Args:
            dist:   (N, *, 1, *)
        Returns:
            (N, *, num_bins, *)
        """
        assert dist.size()[dim] == 1
        offset_shape = [1] * len(dist.size())
        offset_shape[dim] = -1

        if self.use_onehot:
            diff = torch.abs(dist - self.offset.view(*offset_shape))  # (N, *, num_bins, *)
            bin_idx = torch.argmin(diff, dim=dim, keepdim=True)  # (N, *, 1, *)
            y = torch.zeros_like(diff).scatter_(dim=dim, index=bin_idx, value=1.0)
        else:
            overflow_symb = (dist >= self.dist_max).float()  # (N, *, 1, *)
            y = dist - self.offset.view(*offset_shape)  # (N, *, num_bins-1, *)
            y = torch.exp(self.coeff * torch.pow(y, 2))  # (N, *, num_bins-1, *)
            y = torch.cat([y, overflow_symb], dim=dim)  # (N, *, num_bins, *)
            if normalize:
                y = y / y.sum(dim=dim, keepdim=True)

        return y


class PositionalEncoding(nn.Module):
    
    def __init__(self, num_funcs=6):
        super().__init__()
        self.num_funcs = num_funcs
        self.register_buffer('freq_bands', 2.0 ** torch.linspace(0.0, num_funcs-1, num_funcs))
    
    def get_out_dim(self, in_dim):
        return in_dim * (2 * self.num_funcs + 1)

    def forward(self, x):
        """
        Args:
            x:  (..., d).
        """
        shape = list(x.shape[:-1]) + [-1]
        x = x.unsqueeze(-1) # (..., d, 1)
        code = torch.cat([x, torch.sin(x * self.freq_bands), torch.cos(x * self.freq_bands)], dim=-1)   # (..., d, 2f+1)
        code = code.reshape(shape)
        return code


class AngularEncoding(nn.Module):

    def __init__(self, num_funcs=3):
        super().__init__()
        self.num_funcs = num_funcs
        self.register_buffer('freq_bands', torch.FloatTensor(
            [i+1 for i in range(num_funcs)] + [1./(i+1) for i in range(num_funcs)]
        ))

    def get_out_dim(self, in_dim):
        return in_dim * (1 + 2*2*self.num_funcs)

    def forward(self, x):
        """
        Args:
            x:  (..., d).
        """
        shape = list(x.shape[:-1]) + [-1]
        x = x.unsqueeze(-1) # (..., d, 1)
        code = torch.cat([x, torch.sin(x * self.freq_bands), torch.cos(x * self.freq_bands)], dim=-1)   # (..., d, 2f+1)
        code = code.reshape(shape)
        return code


class LayerNorm(nn.Module):

    def __init__(self,
                 normal_shape,
                 gamma=True,
                 beta=True,
                 epsilon=1e-10):
        """Layer normalization layer
        See: [Layer Normalization](https://arxiv.org/pdf/1607.06450.pdf)
        :param normal_shape: The shape of the input tensor or the last dimension of the input tensor.
        :param gamma: Add a scale parameter if it is True.
        :param beta: Add an offset parameter if it is True.
        :param epsilon: Epsilon for calculating variance.
        """
        super().__init__()
        if isinstance(normal_shape, int):
            normal_shape = (normal_shape,)
        else:
            normal_shape = (normal_shape[-1],)
        self.normal_shape = torch.Size(normal_shape)
        self.epsilon = epsilon
        if gamma:
            self.gamma = nn.Parameter(torch.Tensor(*normal_shape))
        else:
            self.register_parameter('gamma', None)
        if beta:
            self.beta = nn.Parameter(torch.Tensor(*normal_shape))
        else:
            self.register_parameter('beta', None)
        self.reset_parameters()

    def reset_parameters(self):
        if self.gamma is not None:
            self.gamma.data.fill_(1)
        if self.beta is not None:
            self.beta.data.zero_()

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
        std = (var + self.epsilon).sqrt()
        y = (x - mean) / std
        if self.gamma is not None:
            y *= self.gamma
        if self.beta is not None:
            y += self.beta
        return y

    def extra_repr(self):
        return 'normal_shape={}, gamma={}, beta={}, epsilon={}'.format(
            self.normal_shape, self.gamma is not None, self.beta is not None, self.epsilon,
        )

AbDesign/diffab/modules/common/layers.py
pengzhangzhi/ab_opt
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .geometry import quaternion_to_rotation_matrix


def log_rotation(R):
    trace = R[..., range(3), range(3)].sum(-1)
    if torch.is_grad_enabled():
        # The derivative of acos at -1.0 is -inf, so to stablize the gradient, we use -0.9999
        min_cos = -0.999
    else:
        min_cos = -1.0
    cos_theta = ( (trace-1) / 2 ).clamp_min(min=min_cos)
    sin_theta = torch.sqrt(1 - cos_theta**2)
    theta = torch.acos(cos_theta)
    coef = ((theta+1e-8)/(2*sin_theta+2e-8))[..., None, None]
    logR = coef * (R - R.transpose(-1, -2))
    return logR


def skewsym_to_so3vec(S):
    x = S[..., 1, 2]
    y = S[..., 2, 0]
    z = S[..., 0, 1]
    w = torch.stack([x,y,z], dim=-1)
    return w


def so3vec_to_skewsym(w):
    x, y, z = torch.unbind(w, dim=-1)
    o = torch.zeros_like(x)
    S = torch.stack([
        o, z, -y,
        -z, o, x,
        y, -x, o,
    ], dim=-1).reshape(w.shape[:-1] + (3, 3))
    return S


def exp_skewsym(S):
    x = torch.linalg.norm(skewsym_to_so3vec(S), dim=-1)
    I = torch.eye(3).to(S).view([1 for _ in range(S.dim()-2)] + [3, 3])
    
    sinx, cosx = torch.sin(x), torch.cos(x)
    b = (sinx + 1e-8) / (x + 1e-8)
    c = (1-cosx + 1e-8) / (x**2 + 2e-8)  # lim_{x->0} (1-cosx)/(x^2) = 0.5

    S2 = S @ S
    return I + b[..., None, None]*S + c[..., None, None]*S2


def so3vec_to_rotation(w):
    return exp_skewsym(so3vec_to_skewsym(w))


def rotation_to_so3vec(R):
    logR = log_rotation(R)
    w = skewsym_to_so3vec(logR)
    return w


def random_uniform_so3(size, device='cpu'):
    q = F.normalize(torch.randn(list(size)+[4,], device=device), dim=-1)    # (..., 4)
    return rotation_to_so3vec(quaternion_to_rotation_matrix(q))


class ApproxAngularDistribution(nn.Module):

    def __init__(self, stddevs, std_threshold=0.1, num_bins=8192, num_iters=1024):
        super().__init__()
        self.std_threshold = std_threshold
        self.num_bins = num_bins
        self.num_iters = num_iters
        self.register_buffer('stddevs', torch.FloatTensor(stddevs))
        self.register_buffer('approx_flag', self.stddevs <= std_threshold)
        self._precompute_histograms()

    @staticmethod
    def _pdf(x, e, L):
        """
        Args:
            x:  (N, )
            e:  Float
            L:  Integer
        """
        x = x[:, None]  # (N, *)
        c = ((1 - torch.cos(x)) / math.pi)  # (N, *)
        l = torch.arange(0, L)[None, :]  # (*, L)
        a = (2*l+1) * torch.exp(-l*(l+1)*(e**2))  # (*, L)
        b = (torch.sin( (l+0.5)* x ) + 1e-6) / (torch.sin( x / 2 ) + 1e-6) # (N, L)
        
        f = (c * a * b).sum(dim=1)
        return f

    def _precompute_histograms(self):
        X, Y = [], []
        for std in self.stddevs:
            std = std.item()
            x = torch.linspace(0, math.pi, self.num_bins)   # (n_bins,)
            y = self._pdf(x, std, self.num_iters)    # (n_bins,)
            y = torch.nan_to_num(y).clamp_min(0)
            X.append(x)
            Y.append(y)
        self.register_buffer('X', torch.stack(X, dim=0))  # (n_stddevs, n_bins)
        self.register_buffer('Y', torch.stack(Y, dim=0))  # (n_stddevs, n_bins)

    def sample(self, std_idx):
        """
        Args:
            std_idx:  Indices of standard deviation.
        Returns:
            samples:  Angular samples [0, PI), same size as std.
        """
        size = std_idx.size()
        std_idx = std_idx.flatten() # (N,)
        
        # Samples from histogram
        prob = self.Y[std_idx]  # (N, n_bins)
        bin_idx = torch.multinomial(prob[:, :-1], num_samples=1).squeeze(-1)    # (N,)
        bin_start = self.X[std_idx, bin_idx]    # (N,)
        bin_width = self.X[std_idx, bin_idx+1] - self.X[std_idx, bin_idx]
        samples_hist = bin_start + torch.rand_like(bin_start) * bin_width    # (N,)

        # Samples from Gaussian approximation
        mean_gaussian = self.stddevs[std_idx]*2
        std_gaussian = self.stddevs[std_idx]
        samples_gaussian = mean_gaussian + torch.randn_like(mean_gaussian) * std_gaussian
        samples_gaussian = samples_gaussian.abs() % math.pi

        # Choose from histogram or Gaussian
        gaussian_flag = self.approx_flag[std_idx]
        samples = torch.where(gaussian_flag, samples_gaussian, samples_hist)

        return samples.reshape(size)


def random_normal_so3(std_idx, angular_distrib, device='cpu'):
    size = std_idx.size()
    u = F.normalize(torch.randn(list(size)+[3,], device=device), dim=-1)
    theta = angular_distrib.sample(std_idx)
    w = u * theta[..., None]
    return w

AbDesign/diffab/modules/common/so3.py
pengzhangzhi/ab_opt
import torch
from torch.nn import Module, Linear, LayerNorm, Sequential, ReLU

from ..common.geometry import compose_rotation_and_translation, quaternion_to_rotation_matrix, repr_6d_to_rotation_matrix


class FrameRotationTranslationPrediction(Module):

    def __init__(self, feat_dim, rot_repr, nn_type='mlp'):
        super().__init__()
        assert rot_repr in ('quaternion', '6d')
        self.rot_repr = rot_repr
        if rot_repr == 'quaternion':
            out_dim = 3 + 3
        elif rot_repr == '6d':
            out_dim = 6 + 3
        
        if nn_type == 'linear':
            self.nn = Linear(feat_dim, out_dim)
        elif nn_type == 'mlp':
            self.nn = Sequential(
                Linear(feat_dim, feat_dim), ReLU(),
                Linear(feat_dim, feat_dim), ReLU(),
                Linear(feat_dim, out_dim)
            )
        else:
            raise ValueError('Unknown nn_type: %s' % nn_type)

    def forward(self, x):
        y = self.nn(x)  # (..., d+3)
        if self.rot_repr == 'quaternion':
            quaternion = torch.cat([torch.ones_like(y[..., :1]), y[..., 0:3]], dim=-1)
            R_delta = quaternion_to_rotation_matrix(quaternion)
            t_delta = y[..., 3:6]
            return R_delta, t_delta
        elif self.rot_repr == '6d':
            R_delta = repr_6d_to_rotation_matrix(y[..., 0:6])
            t_delta = y[..., 6:9]
            return R_delta, t_delta


class FrameUpdate(Module):

    def __init__(self, node_feat_dim, rot_repr='quaternion', rot_tran_nn_type='mlp'):
        super().__init__()
        self.transition_mlp = Sequential(
            Linear(node_feat_dim, node_feat_dim), ReLU(),
            Linear(node_feat_dim, node_feat_dim), ReLU(),
            Linear(node_feat_dim, node_feat_dim),
        )
        self.transition_layer_norm = LayerNorm(node_feat_dim)

        self.rot_tran = FrameRotationTranslationPrediction(node_feat_dim, rot_repr, nn_type=rot_tran_nn_type)
    
    def forward(self, R, t, x, mask_generate):
        """
        Args:
            R:  Frame basis matrices, (N, L, 3, 3_index).
            t:  Frame external (absolute) coordinates, (N, L, 3). Unit: Angstrom.
            x:  Node-wise features, (N, L, F).
            mask_generate:   Masks, (N, L).
        Returns:
            R': Updated basis matrices, (N, L, 3, 3_index).
            t': Updated coordinates, (N, L, 3).
        """
        x = self.transition_layer_norm(x + self.transition_mlp(x))

        R_delta, t_delta = self.rot_tran(x) # (N, L, 3, 3), (N, L, 3)
        R_new, t_new = compose_rotation_and_translation(R, t, R_delta, t_delta)

        mask_R = mask_generate[:, :, None, None].expand_as(R)
        mask_t = mask_generate[:, :, None].expand_as(t)

        R_new = torch.where(mask_R, R_new, R)
        t_new = torch.where(mask_t, t_new, t)

        return R_new, t_new

AbDesign/diffab/modules/common/structure.py
pengzhangzhi/ab_opt
import torch
import torch.nn.functional as F


def get_consecutive_flag(chain_nb, res_nb, mask):
    """
    Args:
        chain_nb, res_nb
    Returns:
        consec: A flag tensor indicating whether residue-i is connected to residue-(i+1), 
                BoolTensor, (B, L-1)[b, i].
    """
    d_res_nb = (res_nb[:, 1:] - res_nb[:, :-1]).abs()   # (B, L-1)
    same_chain = (chain_nb[:, 1:] == chain_nb[:, :-1])
    consec = torch.logical_and(d_res_nb == 1, same_chain)
    consec = torch.logical_and(consec, mask[:, :-1])
    return consec


def get_terminus_flag(chain_nb, res_nb, mask):
    consec = get_consecutive_flag(chain_nb, res_nb, mask)
    N_term_flag = F.pad(torch.logical_not(consec), pad=(1, 0), value=1)
    C_term_flag = F.pad(torch.logical_not(consec), pad=(0, 1), value=1)
    return N_term_flag, C_term_flag

AbDesign/diffab/modules/common/topology.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F
import functools
from tqdm.auto import tqdm

from diffab.modules.common.geometry import apply_rotation_to_vector, quaternion_1ijk_to_rotation_matrix
from diffab.modules.common.so3 import so3vec_to_rotation, rotation_to_so3vec, random_uniform_so3
from diffab.modules.encoders.ga import GAEncoder
from .transition import RotationTransition, PositionTransition, AminoacidCategoricalTransition


def rotation_matrix_cosine_loss(R_pred, R_true):
    """
    Args:
        R_pred: (*, 3, 3).
        R_true: (*, 3, 3).
    Returns:
        Per-matrix losses, (*, ).
    """
    size = list(R_pred.shape[:-2])
    ncol = R_pred.numel() // 3

    RT_pred = R_pred.transpose(-2, -1).reshape(ncol, 3) # (ncol, 3)
    RT_true = R_true.transpose(-2, -1).reshape(ncol, 3) # (ncol, 3)

    ones = torch.ones([ncol, ], dtype=torch.long, device=R_pred.device)
    loss = F.cosine_embedding_loss(RT_pred, RT_true, ones, reduction='none')  # (ncol*3, )
    loss = loss.reshape(size + [3]).sum(dim=-1)    # (*, )
    return loss


class EpsilonNet(nn.Module):

    def __init__(self, res_feat_dim, pair_feat_dim, num_layers, encoder_opt={}):
        super().__init__()
        self.current_sequence_embedding = nn.Embedding(25, res_feat_dim)  # 22 is padding
        self.res_feat_mixer = nn.Sequential(
            nn.Linear(res_feat_dim * 2, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim),
        )
        self.encoder = GAEncoder(res_feat_dim, pair_feat_dim, num_layers, **encoder_opt)

        self.eps_crd_net = nn.Sequential(
            nn.Linear(res_feat_dim+3, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, 3)
        )

        self.eps_rot_net = nn.Sequential(
            nn.Linear(res_feat_dim+3, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, 3)
        )

        self.eps_seq_net = nn.Sequential(
            nn.Linear(res_feat_dim+3, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, 20), nn.Softmax(dim=-1) 
        )

    def forward(self, v_t, p_t, s_t, res_feat, pair_feat, beta, mask_generate, mask_res):
        """
        Args:
            v_t:    (N, L, 3).
            p_t:    (N, L, 3).
            s_t:    (N, L).
            res_feat:   (N, L, res_dim).
            pair_feat:  (N, L, L, pair_dim).
            beta:   (N,).
            mask_generate:    (N, L).
            mask_res:       (N, L).
        Returns:
            v_next: UPDATED (not epsilon) SO3-vector of orietnations, (N, L, 3).
            eps_pos: (N, L, 3).
        """
        N, L = mask_res.size()
        R = so3vec_to_rotation(v_t) # (N, L, 3, 3)

        # s_t = s_t.clamp(min=0, max=19)  # TODO: clamping is good but ugly.
        res_feat = self.res_feat_mixer(torch.cat([res_feat, self.current_sequence_embedding(s_t)], dim=-1)) # [Important] Incorporate sequence at the current step.
        res_feat = self.encoder(R, p_t, res_feat, pair_feat, mask_res)

        t_embed = torch.stack([beta, torch.sin(beta), torch.cos(beta)], dim=-1)[:, None, :].expand(N, L, 3)
        in_feat = torch.cat([res_feat, t_embed], dim=-1)

        # Position changes
        eps_crd = self.eps_crd_net(in_feat)    # (N, L, 3)
        eps_pos = apply_rotation_to_vector(R, eps_crd)  # (N, L, 3)
        eps_pos = torch.where(mask_generate[:, :, None].expand_as(eps_pos), eps_pos, torch.zeros_like(eps_pos))

        # New orientation
        eps_rot = self.eps_rot_net(in_feat)    # (N, L, 3)
        U = quaternion_1ijk_to_rotation_matrix(eps_rot) # (N, L, 3, 3)
        R_next = R @ U
        v_next = rotation_to_so3vec(R_next)     # (N, L, 3)
        v_next = torch.where(mask_generate[:, :, None].expand_as(v_next), v_next, v_t)

        # New sequence categorical distributions
        c_denoised = self.eps_seq_net(in_feat)  # Already softmax-ed, (N, L, 20)

        return v_next, R_next, eps_pos, c_denoised


class FullDPM(nn.Module):

    def __init__(
        self, 
        res_feat_dim, 
        pair_feat_dim, 
        num_steps, 
        eps_net_opt={}, 
        trans_rot_opt={}, 
        trans_pos_opt={}, 
        trans_seq_opt={},
        position_mean=[0.0, 0.0, 0.0],
        position_scale=[10.0],
    ):
        super().__init__()
        self.eps_net = EpsilonNet(res_feat_dim, pair_feat_dim, **eps_net_opt)
        self.num_steps = num_steps
        self.trans_rot = RotationTransition(num_steps, **trans_rot_opt)
        self.trans_pos = PositionTransition(num_steps, **trans_pos_opt)
        self.trans_seq = AminoacidCategoricalTransition(num_steps, **trans_seq_opt)

        self.register_buffer('position_mean', torch.FloatTensor(position_mean).view(1, 1, -1))
        self.register_buffer('position_scale', torch.FloatTensor(position_scale).view(1, 1, -1))
        self.register_buffer('_dummy', torch.empty([0, ]))

    def _normalize_position(self, p):
        p_norm = (p - self.position_mean) / self.position_scale
        return p_norm

    def _unnormalize_position(self, p_norm):
        p = p_norm * self.position_scale + self.position_mean
        return p

    def forward(self, v_0, p_0, s_0, res_feat, pair_feat, mask_generate, mask_res, denoise_structure, denoise_sequence, t=None):
        N, L = res_feat.shape[:2]
        if t == None:
            t = torch.randint(0, self.num_steps, (N,), dtype=torch.long, device=self._dummy.device)
        p_0 = self._normalize_position(p_0)

        if denoise_structure:
            # Add noise to rotation
            R_0 = so3vec_to_rotation(v_0)
            v_noisy, _ = self.trans_rot.add_noise(v_0, mask_generate, t)
            # Add noise to positions
            p_noisy, eps_p = self.trans_pos.add_noise(p_0, mask_generate, t)
        else:
            R_0 = so3vec_to_rotation(v_0)
            v_noisy = v_0.clone()
            p_noisy = p_0.clone()
            eps_p = torch.zeros_like(p_noisy)

        if denoise_sequence:
            # Add noise to sequence
            _, s_noisy = self.trans_seq.add_noise(s_0, mask_generate, t)
        else:
            s_noisy = s_0.clone()

        beta = self.trans_pos.var_sched.betas[t]
        v_pred, R_pred, eps_p_pred, c_denoised = self.eps_net(
            v_noisy, p_noisy, s_noisy, res_feat, pair_feat, beta, mask_generate, mask_res
        )   # (N, L, 3), (N, L, 3, 3), (N, L, 3), (N, L, 20), (N, L)

        loss_dict = {}

        # Rotation loss
        loss_rot = rotation_matrix_cosine_loss(R_pred, R_0) # (N, L)
        loss_rot = (loss_rot * mask_generate).sum() / (mask_generate.sum().float() + 1e-8)
        loss_dict['rot'] = loss_rot

        # Position loss
        loss_pos = F.mse_loss(eps_p_pred, eps_p, reduction='none').sum(dim=-1)  # (N, L)
        loss_pos = (loss_pos * mask_generate).sum() / (mask_generate.sum().float() + 1e-8)
        loss_dict['pos'] = loss_pos

        # Sequence categorical loss
        post_true = self.trans_seq.posterior(s_noisy, s_0, t)
        log_post_pred = torch.log(self.trans_seq.posterior(s_noisy, c_denoised, t) + 1e-8)
        kldiv = F.kl_div(
            input=log_post_pred, 
            target=post_true, 
            reduction='none',
            log_target=False
        ).sum(dim=-1)    # (N, L)
        loss_seq = (kldiv * mask_generate).sum() / (mask_generate.sum().float() + 1e-8)
        loss_dict['seq'] = loss_seq

        return loss_dict

    @torch.no_grad()
    def sample(
        self, 
        v, p, s, 
        res_feat, pair_feat, 
        mask_generate, mask_res, 
        sample_structure=True, sample_sequence=True,
        pbar=False,
    ):
        """
        Args:
            v:  Orientations of contextual residues, (N, L, 3).
            p:  Positions of contextual residues, (N, L, 3).
            s:  Sequence of contextual residues, (N, L).
        """
        N, L = v.shape[:2]
        p = self._normalize_position(p)

        # Set the orientation and position of residues to be predicted to random values
        if sample_structure:
            v_rand = random_uniform_so3([N, L], device=self._dummy.device)
            p_rand = torch.randn_like(p)
            v_init = torch.where(mask_generate[:, :, None].expand_as(v), v_rand, v)
            p_init = torch.where(mask_generate[:, :, None].expand_as(p), p_rand, p)
        else:
            v_init, p_init = v, p

        if sample_sequence:
            s_rand = torch.randint_like(s, low=0, high=19)
            s_init = torch.where(mask_generate, s_rand, s)
        else:
            s_init = s

        traj = {self.num_steps: (v_init, self._unnormalize_position(p_init), s_init)}
        if pbar:
            pbar = functools.partial(tqdm, total=self.num_steps, desc='Sampling')
        else:
            pbar = lambda x: x
        for t in pbar(range(self.num_steps, 0, -1)):
            v_t, p_t, s_t = traj[t]
            p_t = self._normalize_position(p_t)
            
            beta = self.trans_pos.var_sched.betas[t].expand([N, ])
            t_tensor = torch.full([N, ], fill_value=t, dtype=torch.long, device=self._dummy.device)

            v_next, R_next, eps_p, c_denoised = self.eps_net(
                v_t, p_t, s_t, res_feat, pair_feat, beta, mask_generate, mask_res
            )   # (N, L, 3), (N, L, 3, 3), (N, L, 3)

            v_next = self.trans_rot.denoise(v_t, v_next, mask_generate, t_tensor)
            p_next = self.trans_pos.denoise(p_t, eps_p, mask_generate, t_tensor)
            _, s_next = self.trans_seq.denoise(s_t, c_denoised, mask_generate, t_tensor)

            if not sample_structure:
                v_next, p_next = v_t, p_t
            if not sample_sequence:
                s_next = s_t

            traj[t-1] = (v_next, self._unnormalize_position(p_next), s_next)
            traj[t] = tuple(x.cpu() for x in traj[t])    # Move previous states to cpu memory.

        return traj

    @torch.no_grad()
    def optimize(
        self, 
        v, p, s, 
        opt_step: int,
        res_feat, pair_feat, 
        mask_generate, mask_res, 
        sample_structure=True, sample_sequence=True,
        pbar=False,
    ):
        """
        Description:
            First adds noise to the given structure, then denoises it.
        """
        N, L = v.shape[:2]
        p = self._normalize_position(p)
        t = torch.full([N, ], fill_value=opt_step, dtype=torch.long, device=self._dummy.device)

        # Set the orientation and position of residues to be predicted to random values
        if sample_structure:
            # Add noise to rotation
            v_noisy, _ = self.trans_rot.add_noise(v, mask_generate, t)
            # Add noise to positions
            p_noisy, _ = self.trans_pos.add_noise(p, mask_generate, t)
            v_init = torch.where(mask_generate[:, :, None].expand_as(v), v_noisy, v)
            p_init = torch.where(mask_generate[:, :, None].expand_as(p), p_noisy, p)
        else:
            v_init, p_init = v, p

        if sample_sequence:
            _, s_noisy = self.trans_seq.add_noise(s, mask_generate, t)
            s_init = torch.where(mask_generate, s_noisy, s)
        else:
            s_init = s

        traj = {opt_step: (v_init, self._unnormalize_position(p_init), s_init)}
        if pbar:
            pbar = functools.partial(tqdm, total=opt_step, desc='Optimizing')
        else:
            pbar = lambda x: x
        for t in pbar(range(opt_step, 0, -1)):
            v_t, p_t, s_t = traj[t]
            p_t = self._normalize_position(p_t)
            
            beta = self.trans_pos.var_sched.betas[t].expand([N, ])
            t_tensor = torch.full([N, ], fill_value=t, dtype=torch.long, device=self._dummy.device)

            v_next, R_next, eps_p, c_denoised = self.eps_net(
                v_t, p_t, s_t, res_feat, pair_feat, beta, mask_generate, mask_res
            )   # (N, L, 3), (N, L, 3, 3), (N, L, 3)

            v_next = self.trans_rot.denoise(v_t, v_next, mask_generate, t_tensor)
            p_next = self.trans_pos.denoise(p_t, eps_p, mask_generate, t_tensor)
            _, s_next = self.trans_seq.denoise(s_t, c_denoised, mask_generate, t_tensor)

            if not sample_structure:
                v_next, p_next = v_t, p_t
            if not sample_sequence:
                s_next = s_t

            traj[t-1] = (v_next, self._unnormalize_position(p_next), s_next)
            traj[t] = tuple(x.cpu() for x in traj[t])    # Move previous states to cpu memory.

        return traj

AbDesign/diffab/modules/diffusion/dpm_full.py
pengzhangzhi/ab_opt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffab.utils.misc import hotspot_distance_fn, pair2edge,batchfy,clash_loss
from diffab.modules.common.layers import clampped_one_hot
from diffab.modules.common.so3 import ApproxAngularDistribution, random_normal_so3, so3vec_to_rotation, rotation_to_so3vec


class VarianceSchedule(nn.Module):

    def __init__(self, num_steps=100, s=0.01):
        super().__init__()
        T = num_steps
        t = torch.arange(0, num_steps+1, dtype=torch.float)
        f_t = torch.cos( (np.pi / 2) * ((t/T) + s) / (1 + s) ) ** 2
        alpha_bars = f_t / f_t[0]

        betas = 1 - (alpha_bars[1:] / alpha_bars[:-1])
        betas = torch.cat([torch.zeros([1]), betas], dim=0)
        betas = betas.clamp_max(0.999)
 
        sigmas = torch.zeros_like(betas)
        for i in range(1, betas.size(0)):
            sigmas[i] = ((1 - alpha_bars[i-1]) / (1 - alpha_bars[i])) * betas[i]
        sigmas = torch.sqrt(sigmas)
        
        self.register_buffer('betas', betas)
        self.register_buffer('alpha_bars', alpha_bars)
        self.register_buffer('alphas', 1 - betas)
        self.register_buffer('sigmas', sigmas)
        # calculate X0 = sqrt_recip_alphas_cumprod * Xt - sqrt_recipm1_alphas_cumprod * noise
        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alpha_bars))
        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alpha_bars - 1))
        
        alphas = 1 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
        self.posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
        self.posterior_log_variance_clipped = torch.log(self.posterior_variance.clamp(min =1e-20))
        self.posterior_mean_coef1 = betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)
        self.posterior_mean_coef2 = (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod)
        
    def to(self,device):
        for k, v in self.__dict__.items():
            if isinstance(v, torch.Tensor):
                self.__dict__[k] = v.to(device)
                
        
        
class PositionTransition(nn.Module):

    def __init__(self, num_steps, var_sched_opt={}):
        super().__init__()
        self.var_sched = VarianceSchedule(num_steps, **var_sched_opt)
         
    def add_noise(self, p_0, mask_generate,  t):
        """
        Args:
            p_0:    (N, L, 3).
            mask_generate:    (N, L).
            t:  (N,).
        """
        alpha_bar = self.var_sched.alpha_bars[t]
        nb_non_batch_dims = len(p_0.shape) - 1
        c0 = torch.sqrt(alpha_bar).view(-1, *(nb_non_batch_dims*(1,)))
        c1 = torch.sqrt(1 - alpha_bar).view(-1, *(nb_non_batch_dims*(1,)))
        prior_p = torch.randn_like(p_0)
           
        p_noisy = c0*p_0 + c1*prior_p
        p_noisy = torch.where(mask_generate[..., None].expand_as(p_0), p_noisy, p_0)

        return p_noisy, prior_p
    
    def pred_start_from_noise(self, p_t, eps_p, mask_generate, t):
        nb_non_batch_dims = len(p_t.shape) - 1
        sqrt_recip_alphas_cumprod = self.var_sched.sqrt_recip_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        sqrt_recipm1_alphas_cumprod = self.var_sched.sqrt_recipm1_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        
        p_0 = sqrt_recip_alphas_cumprod * p_t - sqrt_recipm1_alphas_cumprod * eps_p

        p_0 = torch.where(mask_generate[..., None].expand_as(p_t), p_0, p_t)
        return p_0
    
    def pred_noise_from_start(self, p_t, p_0, mask_generate, t):
        nb_non_batch_dims = len(p_t.shape) - 1
        sqrt_recip_alphas_cumprod = self.var_sched.sqrt_recip_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        sqrt_recipm1_alphas_cumprod = self.var_sched.sqrt_recipm1_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        eps_p = (sqrt_recip_alphas_cumprod * p_t - p_0) / sqrt_recipm1_alphas_cumprod

        eps_p = torch.where(mask_generate[..., None].expand_as(p_t), eps_p, p_t)
        return eps_p
    def denoise_from_p0(self, p_t, p_0, mask_generate, t,guidance_kwargs=None):
        self.var_sched.to(p_t.device)
        nb_non_batch_dims = len(p_t.shape) - 1
        posterior_mean = (
            self.var_sched.posterior_mean_coef1[t].view(-1, *(nb_non_batch_dims*(1,))) * p_0 +
            self.var_sched.posterior_mean_coef2[t].view(-1, *(nb_non_batch_dims*(1,))) * p_t
        )
        posterior_variance = self.var_sched.posterior_variance[t].view(-1, *(nb_non_batch_dims*(1,)))
        posterior_log_variance_clipped = self.var_sched.posterior_log_variance_clipped[t].view(-1, *(nb_non_batch_dims*(1,)))
        
        # with torch.enable_grad():
        #     batched_mean = batchfy(posterior_mean,guidance_kwargs['lengths']).detach().requires_grad_(True)
        #     hotspot_mask = guidance_kwargs['to_hotspot_dist'] != 0
        #     hotspot_idx = (torch.where(hotspot_mask)[1])
        #     cdr_idx = (torch.where(guidance_kwargs['generate_flag'])[1])
        #     cdr2hotspot_dist = hotspot_distance_fn(batched_mean, hotspot_idx, cdr_idx)
            
        #     grad = torch.autograd.grad(cdr2hotspot_dist.sum(), batched_mean)[0] * 1
        #     grad = grad[guidance_kwargs['mask']]
        #     posterior_mean = posterior_mean -  grad #* posterior_variance 
        #     print(f"\n dist: {10*cdr2hotspot_dist.mean():.3f},\
        #          native dist: {guidance_kwargs['to_hotspot_dist'].max():.3f}\
        #         grad: {grad[mask_generate].mean()},  posterior_mean {posterior_mean[mask_generate].mean()}")  
        
        z = torch.where(
            (t > 1)[...,None].expand_as(p_t),
            torch.randn_like(p_t),
            torch.zeros_like(p_t),
        )
        p_next = posterior_mean + (0.5 * posterior_log_variance_clipped).exp() * z
        p_next = torch.where(mask_generate[..., None].expand_as(p_t), p_next, p_t)
        
        return p_next
    
    
    def denoise(self, p_t, eps_p, mask_generate, t):
        # IMPORTANT:
        #   clampping alpha is to fix the instability issue at the first step (t=T)
        #   it seems like a problem with the ``improved ddpm''.
        alpha = self.var_sched.alphas[t].clamp_min(
            self.var_sched.alphas[-2]
        )
        nb_non_batch_dims = len(p_t.shape) - 1
        alpha_bar = self.var_sched.alpha_bars[t]
        sigma = self.var_sched.sigmas[t].view(-1, *(nb_non_batch_dims*(1,)))

        c0 = ( 1.0 / torch.sqrt(alpha + 1e-8) ).view(-1, *(nb_non_batch_dims*(1,)))
        c1 = ( (1 - alpha) / torch.sqrt(1 - alpha_bar + 1e-8) ).view(-1, *(nb_non_batch_dims*(1,)))
        
        if len(mask_generate.shape) == 2:
            t = t[:, None].expand(*mask_generate.shape)
        elif len(mask_generate.shape) == 1:
            assert len(t.shape) == 1
        else:
            raise ValueError('mask_generate must be 1D or 2D')
        
        z = torch.where(
            (t > 1)[...,None].expand_as(p_t),
            torch.randn_like(p_t),
            torch.zeros_like(p_t),
        )
        mean = c0 * (p_t - c1 * eps_p)
        p_next = mean + sigma * z
        p_next = torch.where(mask_generate[..., None].expand_as(p_t), p_next, p_t)
        return p_next


class RotationTransition(nn.Module):

    def __init__(self, num_steps, var_sched_opt={}, angular_distrib_fwd_opt={}, angular_distrib_inv_opt={}):
        super().__init__()
        self.var_sched = VarianceSchedule(num_steps, **var_sched_opt)

        # Forward (perturb)
        c1 = torch.sqrt(1 - self.var_sched.alpha_bars) # (T,).
        self.angular_distrib_fwd = ApproxAngularDistribution(c1.tolist(), **angular_distrib_fwd_opt)

        # Inverse (generate)
        sigma = self.var_sched.sigmas
        self.angular_distrib_inv = ApproxAngularDistribution(sigma.tolist(), **angular_distrib_inv_opt)

        self.register_buffer('_dummy', torch.empty([0, ]))
    
    def add_noise(self, v_0, mask_generate, t, ):
        """
        Args:
            v_0:    (N, L, 3).
            mask_generate:    (N, L).
            t:  (N,).
        """
        nb_non_batch_dims = len(v_0.shape) - 1
        alpha_bar = self.var_sched.alpha_bars[t]
        c0 = torch.sqrt(alpha_bar).view(-1, *(nb_non_batch_dims*(1,)))
        c1 = torch.sqrt(1 - alpha_bar).view(-1, *(nb_non_batch_dims*(1,)))

        # Noise rotation
        if len(mask_generate.shape) == 2:
            t = t[:, None].expand(*mask_generate.shape)
        elif len(mask_generate.shape) == 1:
            assert len(t.shape) == 1
        else:
            raise ValueError('mask_generate must be 1D or 2D')
        e_scaled = random_normal_so3(t, self.angular_distrib_fwd, device=self._dummy.device)    # (N, L, 3)
        e_normal = e_scaled / (c1 + 1e-8)
        E_scaled = so3vec_to_rotation(e_scaled)   # (N, L, 3, 3)

        # Scaled true rotation
        R0_scaled = so3vec_to_rotation(c0 * v_0)  # (N, L, 3, 3)

        R_noisy = E_scaled @ R0_scaled
        v_noisy = rotation_to_so3vec(R_noisy)
        v_noisy = torch.where(mask_generate[..., None].expand_as(v_0), v_noisy, v_0)

        return v_noisy, e_scaled

    def denoise(self, v_t, v_next, mask_generate, t):
        if len(mask_generate.shape) == 2:
            # [N, L]
            t = t[:, None].expand(*mask_generate.shape)
        elif len(mask_generate.shape) == 1:
            # [nb_all_res,]
            assert len(t.shape) == 1
        else:
            raise ValueError('mask_generate must be 1D or 2D')
        
        e = random_normal_so3(t, self.angular_distrib_inv, device=self._dummy.device) # (N, L, 3)
        e = torch.where(
            (t > 1)[...,None].expand_as(e),
            e, 
            torch.zeros_like(e) # Simply denoise and don't add noise at the last step
        )
        E = so3vec_to_rotation(e)

        R_next = E @ so3vec_to_rotation(v_next)
        v_next = rotation_to_so3vec(R_next)
        v_next = torch.where(mask_generate[..., None].expand_as(v_next), v_next, v_t)

        return v_next


class AminoacidCategoricalTransition(nn.Module):
    
    def __init__(self, num_steps, num_classes=20, var_sched_opt={}):
        super().__init__()
        self.num_classes = num_classes
        self.var_sched = VarianceSchedule(num_steps, **var_sched_opt)

    @staticmethod
    def _sample(c):
        """
        Args:
            c:    (N, L, K).
        Returns:
            x:    (N, L).
        """
        N, L, K = c.size()
        c = c.view(N*L, K) + 1e-8
        x = torch.multinomial(c, 1).view(N, L)
        return x

    def add_noise(self, x_0, mask_generate, t):
        """
        Args:
            x_0:    (N, L)
            mask_generate:    (N, L).
            t:  (N,).
        Returns:
            c_t:    Probability, (N, L, K).
            x_t:    Sample, LongTensor, (N, L).
        """
        N, L = x_0.size()
        K = self.num_classes
        c_0 = clampped_one_hot(x_0, num_classes=K).float() # (N, L, K).
        alpha_bar = self.var_sched.alpha_bars[t][:, None, None] # (N, 1, 1)
        c_noisy = (alpha_bar*c_0) + ( (1-alpha_bar)/K )
        c_t = torch.where(mask_generate[..., None].expand(N,L,K), c_noisy, c_0)
        x_t = self._sample(c_t)
        return c_t, x_t

    def posterior(self, x_t, x_0, t):
        """
        Args:
            x_t:    Category LongTensor (N, L) or Probability FloatTensor (N, L, K).
            x_0:    Category LongTensor (N, L) or Probability FloatTensor (N, L, K).
            t:  (N,).
        Returns:
            theta:  Posterior probability at (t-1)-th step, (N, L, K).
        """
        K = self.num_classes

        if x_t.dim() == 3:
            c_t = x_t   # When x_t is probability distribution.
        else:
            c_t = clampped_one_hot(x_t, num_classes=K).float() # (N, L, K)

        if x_0.dim() == 3:
            c_0 = x_0   # When x_0 is probability distribution.
        else:
            c_0 = clampped_one_hot(x_0, num_classes=K).float() # (N, L, K)

        alpha = self.var_sched.alpha_bars[t][:, None, None]     # (N, 1, 1)
        alpha_bar = self.var_sched.alpha_bars[t][:, None, None] # (N, 1, 1)

        theta = ((alpha*c_t) + (1-alpha)/K) * ((alpha_bar*c_0) + (1-alpha_bar)/K)   # (N, L, K)
        theta = theta / (theta.sum(dim=-1, keepdim=True) + 1e-8)
        return theta

    def denoise(self, x_t, c_0_pred, mask_generate, t):
        """
        Args:
            x_t:        (N, L).
            c_0_pred:   Normalized probability predicted by networks, (N, L, K).
            mask_generate:    (N, L).
            t:  (N,).
        Returns:
            post:   Posterior probability at (t-1)-th step, (N, L, K).
            x_next: Sample at (t-1)-th step, LongTensor, (N, L).
        """
        c_t = clampped_one_hot(x_t, num_classes=self.num_classes).float()  # (N, L, K)
        post = self.posterior(c_t, c_0_pred, t=t)   # (N, L, K)
        post = torch.where(mask_generate[..., None].expand(post.size()), post, c_t)
        x_next = self._sample(post)
        return post, x_next

AbDesign/diffab/modules/diffusion/transition.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from diffab.modules.common.geometry import global_to_local, local_to_global, normalize_vector, construct_3d_basis, angstrom_to_nm
from diffab.modules.common.layers import mask_zero, LayerNorm
from diffab.utils.protein.constants import BBHeavyAtom


def _alpha_from_logits(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(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)


class GABlock(nn.Module):

    def __init__(self, node_feat_dim, pair_feat_dim, value_dim=32, query_key_dim=32, num_query_points=8,
                 num_value_points=8, num_heads=12, bias=False):
        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=bias)
        self.proj_key = nn.Linear(node_feat_dim, query_key_dim * num_heads, bias=bias)
        self.proj_value = nn.Linear(node_feat_dim, value_dim * num_heads, bias=bias)

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

        # Spatial
        self.spatial_coef = nn.Parameter(torch.full([1, 1, 1, self.num_heads], fill_value=np.log(np.exp(1.) - 1.)),
                                         requires_grad=True)
        self.proj_query_point = nn.Linear(node_feat_dim, num_query_points * num_heads * 3, bias=bias)
        self.proj_key_point = nn.Linear(node_feat_dim, num_query_points * num_heads * 3, bias=bias)
        self.proj_value_point = nn.Linear(node_feat_dim, num_value_points * num_heads * 3, bias=bias)

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

        self.layer_norm_1 = LayerNorm(node_feat_dim)
        self.mlp_transition = nn.Sequential(nn.Linear(node_feat_dim, node_feat_dim), nn.ReLU(),
                                            nn.Linear(node_feat_dim, node_feat_dim), nn.ReLU(),
                                            nn.Linear(node_feat_dim, node_feat_dim))
        self.layer_norm_2 = LayerNorm(node_feat_dim)

    def _node_logits(self, x):
        query_l = _heads(self.proj_query(x), self.num_heads, self.query_key_dim)  # (N, L, n_heads, qk_ch)
        key_l = _heads(self.proj_key(x), self.num_heads, self.query_key_dim)  # (N, L, n_heads, qk_ch)
        logits_node = (query_l.unsqueeze(2) * key_l.unsqueeze(1) *
                       (1 / np.sqrt(self.query_key_dim))).sum(-1)  # (N, L, L, num_heads)
        return logits_node

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

    def _spatial_logits(self, R, t, x):
        N, L, _ = t.size()

        # Query
        query_points = _heads(self.proj_query_point(x), self.num_heads * self.num_query_points,
                              3)  # (N, L, n_heads * n_pnts, 3)
        query_points = local_to_global(R, t, query_points)  # Global query coordinates, (N, L, n_heads * n_pnts, 3)
        query_s = query_points.reshape(N, L, self.num_heads, -1)  # (N, L, n_heads, n_pnts*3)

        # Key
        key_points = _heads(self.proj_key_point(x), self.num_heads * self.num_query_points,
                            3)  # (N, L, 3, n_heads * n_pnts)
        key_points = local_to_global(R, t, key_points)  # Global key coordinates, (N, L, n_heads * n_pnts, 3)
        key_s = key_points.reshape(N, L, self.num_heads, -1)  # (N, L, n_heads, n_pnts*3)

        # Q-K Product
        sum_sq_dist = ((query_s.unsqueeze(2) - key_s.unsqueeze(1)) ** 2).sum(-1)  # (N, L, L, n_heads)
        gamma = F.softplus(self.spatial_coef)
        logits_spatial = sum_sq_dist * ((-1 * gamma * np.sqrt(2 / (9 * self.num_query_points)))
                                        / 2)  # (N, L, L, n_heads)
        return logits_spatial

    def _pair_aggregation(self, alpha, z):
        N, L = z.shape[:2]
        feat_p2n = alpha.unsqueeze(-1) * z.unsqueeze(-2)  # (N, L, L, n_heads, C)
        feat_p2n = feat_p2n.sum(dim=2)  # (N, L, n_heads, C)
        return feat_p2n.reshape(N, L, -1)

    def _node_aggregation(self, alpha, x):
        N, L = x.shape[:2]
        value_l = _heads(self.proj_value(x), self.num_heads, self.query_key_dim)  # (N, L, n_heads, v_ch)
        feat_node = alpha.unsqueeze(-1) * value_l.unsqueeze(1)  # (N, L, L, n_heads, *) @ (N, *, L, n_heads, v_ch)
        feat_node = feat_node.sum(dim=2)  # (N, L, n_heads, v_ch)
        return feat_node.reshape(N, L, -1)

    def _spatial_aggregation(self, alpha, R, t, x):
        N, L, _ = t.size()
        value_points = _heads(self.proj_value_point(x), self.num_heads * self.num_value_points,
                              3)  # (N, L, n_heads * n_v_pnts, 3)
        value_points = local_to_global(R, t, value_points.reshape(N, L, self.num_heads, self.num_value_points,
                                                                  3))  # (N, L, n_heads, n_v_pnts, 3)
        aggr_points = alpha.reshape(N, L, L, self.num_heads, 1, 1) * \
                      value_points.unsqueeze(1)  # (N, *, L, n_heads, n_pnts, 3)
        aggr_points = aggr_points.sum(dim=2)  # (N, L, n_heads, n_pnts, 3)

        feat_points = global_to_local(R, t, aggr_points)  # (N, L, n_heads, n_pnts, 3)
        feat_distance = feat_points.norm(dim=-1)  # (N, L, n_heads, n_pnts)
        feat_direction = normalize_vector(feat_points, dim=-1, eps=1e-4)  # (N, L, n_heads, n_pnts, 3)

        feat_spatial = torch.cat([
            feat_points.reshape(N, L, -1),
            feat_distance.reshape(N, L, -1),
            feat_direction.reshape(N, L, -1),
        ], dim=-1)

        return feat_spatial

    def forward(self, R, t, x, z, mask):
        """
        Args:
            R:  Frame basis matrices, (N, L, 3, 3_index).
            t:  Frame external (absolute) coordinates, (N, L, 3).
            x:  Node-wise features, (N, L, F).
            z:  Pair-wise features, (N, L, L, C).
            mask:   Masks, (N, L).
        Returns:
            x': Updated node-wise features, (N, L, F).
        """
        # Attention logits
        logits_node = self._node_logits(x)
        logits_pair = self._pair_logits(z)
        logits_spatial = self._spatial_logits(R, t, x)
        # Summing logits up and apply `softmax`.
        logits_sum = logits_node + logits_pair + logits_spatial
        alpha = _alpha_from_logits(logits_sum * np.sqrt(1 / 3), mask)  # (N, L, L, n_heads)

        # Aggregate features
        feat_p2n = self._pair_aggregation(alpha, z)
        feat_node = self._node_aggregation(alpha, x)
        feat_spatial = self._spatial_aggregation(alpha, R, t, x)

        # Finally
        feat_all = self.out_transform(torch.cat([feat_p2n, feat_node, feat_spatial], dim=-1))  # (N, L, F)
        feat_all = mask_zero(mask.unsqueeze(-1), feat_all)
        x_updated = self.layer_norm_1(x + feat_all)
        x_updated = self.layer_norm_2(x_updated + self.mlp_transition(x_updated))
        return x_updated


class GAEncoder(nn.Module):

    def __init__(self, node_feat_dim, pair_feat_dim, num_layers, ga_block_opt={}):
        super(GAEncoder, self).__init__()
        self.blocks = nn.ModuleList([
            GABlock(node_feat_dim, pair_feat_dim, **ga_block_opt) 
            for _ in range(num_layers)
        ])

    def forward(self, R, t, res_feat, pair_feat, mask):
        for i, block in enumerate(self.blocks):
            res_feat = block(R, t, res_feat, pair_feat, mask)
        return res_feat

AbDesign/diffab/modules/encoders/ga.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F

from diffab.modules.common.geometry import angstrom_to_nm, pairwise_dihedrals
from diffab.modules.common.layers import AngularEncoding
from diffab.utils.protein.constants import BBHeavyAtom, AA


class PairEmbedding(nn.Module):

    def __init__(self, feat_dim, max_num_atoms, max_aa_types=22, max_relpos=32):
        super().__init__()
        self.max_num_atoms = max_num_atoms
        self.max_aa_types = max_aa_types
        self.max_relpos = max_relpos
        self.aa_pair_embed = nn.Embedding(self.max_aa_types*self.max_aa_types, feat_dim)
        self.relpos_embed = nn.Embedding(2*max_relpos+1, feat_dim)

        self.aapair_to_distcoef = nn.Embedding(self.max_aa_types*self.max_aa_types, max_num_atoms*max_num_atoms)
        nn.init.zeros_(self.aapair_to_distcoef.weight)
        self.distance_embed = nn.Sequential(
            nn.Linear(max_num_atoms*max_num_atoms, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim), nn.ReLU(),
        )

        self.dihedral_embed = AngularEncoding()
        feat_dihed_dim = self.dihedral_embed.get_out_dim(2) # Phi and Psi

        infeat_dim = feat_dim+feat_dim+feat_dim+feat_dihed_dim
        self.out_mlp = nn.Sequential(
            nn.Linear(infeat_dim, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim),
        )

    def forward(self, aa, res_nb, chain_nb, pos_atoms, mask_atoms, structure_mask=None, sequence_mask=None):
        """
        Args:
            aa: (N, L).
            res_nb: (N, L).
            chain_nb: (N, L).
            pos_atoms:  (N, L, A, 3)
            mask_atoms: (N, L, A)
            structure_mask: (N, L)
            sequence_mask:  (N, L), mask out unknown amino acids to generate.

        Returns:
            (N, L, L, feat_dim)
        """
        N, L = aa.size()

        # Remove other atoms
        pos_atoms = pos_atoms[:, :, :self.max_num_atoms]
        mask_atoms = mask_atoms[:, :, :self.max_num_atoms]

        mask_residue = mask_atoms[:, :, BBHeavyAtom.CA] # (N, L)
        mask_pair = mask_residue[:, :, None] * mask_residue[:, None, :]
        pair_structure_mask = structure_mask[:, :, None] * structure_mask[:, None, :] if structure_mask is not None else None

        # Pair identities
        if sequence_mask is not None:
            # Avoid data leakage at training time
            aa = torch.where(sequence_mask, aa, torch.full_like(aa, fill_value=AA.UNK))
        aa_pair = aa[:,:,None]*self.max_aa_types + aa[:,None,:]    # (N, L, L)
        feat_aapair = self.aa_pair_embed(aa_pair)
    
        # Relative sequential positions
        same_chain = (chain_nb[:, :, None] == chain_nb[:, None, :])
        relpos = torch.clamp(
            res_nb[:,:,None] - res_nb[:,None,:], 
            min=-self.max_relpos, max=self.max_relpos,
        )   # (N, L, L)
        feat_relpos = self.relpos_embed(relpos + self.max_relpos) * same_chain[:,:,:,None]

        # Distances
        d = angstrom_to_nm(torch.linalg.norm(
            pos_atoms[:,:,None,:,None] - pos_atoms[:,None,:,None,:],
            dim = -1, ord = 2,
        )).reshape(N, L, L, -1) # (N, L, L, A*A)
        c = F.softplus(self.aapair_to_distcoef(aa_pair))    # (N, L, L, A*A)
        d_gauss = torch.exp(-1 * c * d**2)
        mask_atom_pair = (mask_atoms[:,:,None,:,None] * mask_atoms[:,None,:,None,:]).reshape(N, L, L, -1)
        feat_dist = self.distance_embed(d_gauss * mask_atom_pair)
        if pair_structure_mask is not None:
            # Avoid data leakage at training time
            feat_dist = feat_dist * pair_structure_mask[:, :, :, None]

        # Orientations
        dihed = pairwise_dihedrals(pos_atoms)   # (N, L, L, 2)
        feat_dihed = self.dihedral_embed(dihed)
        if pair_structure_mask is not None:
            # Avoid data leakage at training time
            feat_dihed = feat_dihed * pair_structure_mask[:, :, :, None]

        # All
        feat_all = torch.cat([feat_aapair, feat_relpos, feat_dist, feat_dihed], dim=-1)
        feat_all = self.out_mlp(feat_all)   # (N, L, L, F)
        feat_all = feat_all * mask_pair[:, :, :, None]

        return feat_all


AbDesign/diffab/modules/encoders/pair.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn

from diffab.modules.common.geometry import construct_3d_basis, global_to_local, get_backbone_dihedral_angles
from diffab.modules.common.layers import AngularEncoding
from diffab.utils.protein.constants import BBHeavyAtom, AA


class ResidueEmbedding(nn.Module):

    def __init__(self, feat_dim, max_num_atoms, max_aa_types=22,hotspot=True):
        super().__init__()
        self.max_num_atoms = max_num_atoms
        self.max_aa_types = max_aa_types
        self.aatype_embed = nn.Embedding(self.max_aa_types, feat_dim)
        self.dihed_embed = AngularEncoding()
        self.type_embed = nn.Embedding(10, feat_dim, padding_idx=0)    # 1: Heavy, 2: Light, 3: Ag
        infeat_dim = feat_dim + (self.max_aa_types*max_num_atoms*3) + self.dihed_embed.get_out_dim(3) + feat_dim
        if hotspot:
            infeat_dim += feat_dim
            self.hotspot_embed = nn.Embedding(10, feat_dim, padding_idx=0)    # 1: non-hotspot, 2: hotspot 
        self.mlp = nn.Sequential(
            nn.Linear(infeat_dim, feat_dim * 2), nn.ReLU(),
            nn.Linear(feat_dim * 2, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim)
        )

    def forward(
        self,
        aa, res_nb, chain_nb, pos_atoms, mask_atoms, fragment_type, 
        hotspot=None, hydropathy=None, charge=None, 
        structure_mask=None, sequence_mask=None
        ):
        """
        Args:
            aa:         (N, L).
            res_nb:     (N, L).
            chain_nb:   (N, L).
            pos_atoms:  (N, L, A, 3).
            mask_atoms: (N, L, A).
            fragment_type:  (N, L).
            structure_mask: (N, L), mask out unknown structures to generate.
            sequence_mask:  (N, L), mask out unknown amino acids to generate.
        """
        N, L = aa.size()
        mask_residue = mask_atoms[:, :, BBHeavyAtom.CA] # (N, L)

        # Remove other atoms
        pos_atoms = pos_atoms[:, :, :self.max_num_atoms]
        mask_atoms = mask_atoms[:, :, :self.max_num_atoms]

        # Amino acid identity features
        if sequence_mask is not None:
            # Avoid data leakage at training time
            aa = torch.where(sequence_mask, aa, torch.full_like(aa, fill_value=AA.UNK))
        aa_feat = self.aatype_embed(aa) # (N, L, feat)

        # Coordinate features
        R = construct_3d_basis(
            pos_atoms[:, :, BBHeavyAtom.CA], 
            pos_atoms[:, :, BBHeavyAtom.C], 
            pos_atoms[:, :, BBHeavyAtom.N]
        )
        t = pos_atoms[:, :, BBHeavyAtom.CA]
        crd = global_to_local(R, t, pos_atoms)    # (N, L, A, 3)
        crd_mask = mask_atoms[:, :, :, None].expand_as(crd)
        crd = torch.where(crd_mask, crd, torch.zeros_like(crd))

        aa_expand  = aa[:, :, None, None, None].expand(N, L, self.max_aa_types, self.max_num_atoms, 3)
        rng_expand = torch.arange(0, self.max_aa_types)[None, None, :, None, None].expand(N, L, self.max_aa_types, self.max_num_atoms, 3).to(aa_expand)
        place_mask = (aa_expand == rng_expand)
        crd_expand = crd[:, :, None, :, :].expand(N, L, self.max_aa_types, self.max_num_atoms, 3)
        crd_expand = torch.where(place_mask, crd_expand, torch.zeros_like(crd_expand))
        crd_feat = crd_expand.reshape(N, L, self.max_aa_types*self.max_num_atoms*3)
        if structure_mask is not None:
            # Avoid data leakage at training time
            crd_feat = crd_feat * structure_mask[:, :, None]

        # Backbone dihedral features
        bb_dihedral, mask_bb_dihed = get_backbone_dihedral_angles(pos_atoms, chain_nb=chain_nb, res_nb=res_nb, mask=mask_residue)
        dihed_feat = self.dihed_embed(bb_dihedral[:, :, :, None]) * mask_bb_dihed[:, :, :, None]  # (N, L, 3, dihed/3)
        dihed_feat = dihed_feat.reshape(N, L, -1)
        if structure_mask is not None:
            # Avoid data leakage at training time
            dihed_mask = torch.logical_and(
                structure_mask,
                torch.logical_and(
                    torch.roll(structure_mask, shifts=+1, dims=1), 
                    torch.roll(structure_mask, shifts=-1, dims=1)
                ),
            )   # Avoid slight data leakage via dihedral angles of anchor residues
            dihed_feat = dihed_feat * dihed_mask[:, :, None]

        # Type feature
        type_feat = self.type_embed(fragment_type) # (N, L, feat)
        out = torch.cat([aa_feat, crd_feat, dihed_feat, type_feat], dim=-1)
        if hotspot is not None and self.hotspot_embed is not None:
            hotspot_feat = self.hotspot_embed(hotspot)
            out = torch.cat([out, hotspot_feat], dim=-1)
        out_feat = self.mlp(out) # (N, L, F)
        out_feat = out_feat * mask_residue[:, :, None]
        return out_feat

AbDesign/diffab/modules/encoders/residue.py
pengzhangzhi/ab_opt
import abc
from typing import List


FilePath = str


class DockingEngine(abc.ABC):

    @abc.abstractmethod
    def __enter__(self):
        pass

    @abc.abstractmethod
    def __exit__(self, typ, value, traceback):
        pass

    @abc.abstractmethod
    def set_receptor(self, pdb_path: FilePath):
        pass

    @abc.abstractmethod
    def set_ligand(self, pdb_path: FilePath):
        pass

    @abc.abstractmethod
    def dock(self) -> List[FilePath]:
        pass

AbDesign/diffab/tools/dock/base.py
pengzhangzhi/ab_opt
import os
import shutil
import tempfile
import subprocess
import dataclasses as dc
from typing import List, Optional
from Bio import PDB
from Bio.PDB import Model as PDBModel

from diffab.tools.renumber import renumber as renumber_chothia
from .base import DockingEngine


def fix_docked_pdb(pdb_path):
    fixed = []
    with open(pdb_path, 'r') as f:
        for ln in f.readlines():
            if (ln.startswith('ATOM') or ln.startswith('HETATM')) and len(ln) == 56:
                fixed.append( ln[:-1] + ' 1.00  0.00              \n' )
            else:
                fixed.append(ln)
    with open(pdb_path, 'w') as f:
        f.write(''.join(fixed))


class HDock(DockingEngine):

    def __init__(
        self, 
        hdock_bin='./bin/hdock',
        createpl_bin='./bin/createpl',
    ):
        super().__init__()
        self.hdock_bin = os.path.realpath(hdock_bin)
        self.createpl_bin = os.path.realpath(createpl_bin)
        self.tmpdir = tempfile.TemporaryDirectory()

        self._has_receptor = False
        self._has_ligand = False

        self._receptor_chains = []
        self._ligand_chains = []

    def __enter__(self):
        return self

    def __exit__(self, typ, value, traceback):
        self.tmpdir.cleanup()

    def set_receptor(self, pdb_path):
        shutil.copyfile(pdb_path, os.path.join(self.tmpdir.name, 'receptor.pdb'))
        self._has_receptor = True

    def set_ligand(self, pdb_path):
        shutil.copyfile(pdb_path, os.path.join(self.tmpdir.name, 'ligand.pdb'))
        self._has_ligand = True

    def _dump_complex_pdb(self):
        parser = PDB.PDBParser(QUIET=True)
        model_receptor = parser.get_structure(None, os.path.join(self.tmpdir.name, 'receptor.pdb'))[0]
        docked_pdb_path = os.path.join(self.tmpdir.name, 'ligand_docked.pdb')
        fix_docked_pdb(docked_pdb_path)
        structure_ligdocked = parser.get_structure(None, docked_pdb_path)

        pdb_io = PDB.PDBIO()
        paths = []
        for i, model_ligdocked in enumerate(structure_ligdocked):
            model_complex = PDBModel.Model(0)
            for chain in model_receptor:
                model_complex.add(chain.copy())
            for chain in model_ligdocked:
                model_complex.add(chain.copy())
            pdb_io.set_structure(model_complex)
            save_path = os.path.join(self.tmpdir.name, f"complex_{i}.pdb")
            pdb_io.save(save_path)
            paths.append(save_path)
        return paths

    def dock(self):
        if not (self._has_receptor and self._has_ligand):
            raise ValueError('Missing receptor or ligand.')
        subprocess.run(
            [self.hdock_bin, "receptor.pdb", "ligand.pdb"],
            cwd=self.tmpdir.name, check=True
        )
        subprocess.run(
            [self.createpl_bin, "Hdock.out", "ligand_docked.pdb"], 
            cwd=self.tmpdir.name, check=True
        )
        return self._dump_complex_pdb()


@dc.dataclass
class DockSite:
    chain: str
    resseq: int


class HDockAntibody(HDock):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._heavy_chain_id = None
        self._epitope_sites: Optional[List[DockSite]] = None

    def set_ligand(self, pdb_path):
        raise NotImplementedError('Please use set_antibody')
    
    def set_receptor(self, pdb_path):
        raise NotImplementedError('Please use set_antigen')

    def set_antigen(self, pdb_path, epitope_sites: Optional[List[DockSite]]=None):
        super().set_receptor(pdb_path)
        self._epitope_sites = epitope_sites

    def set_antibody(self, pdb_path):
        heavy_chains, _ = renumber_chothia(pdb_path, os.path.join(self.tmpdir.name, 'ligand.pdb'))
        self._has_ligand = True
        self._heavy_chain_id = heavy_chains[0]

    def _prepare_lsite(self):
        lsite_content = f"95-102:{self._heavy_chain_id}\n"  # Chothia CDR H3
        with open(os.path.join(self.tmpdir.name, 'lsite.txt'), 'w') as f:
            f.write(lsite_content)
        print(f"[INFO] lsite content: {lsite_content}")

    def _prepare_rsite(self):
        rsite_content = ""
        for site in self._epitope_sites:
            rsite_content += f"{site.resseq}:{site.chain}\n"
        with open(os.path.join(self.tmpdir.name, 'rsite.txt'), 'w') as f:
            f.write(rsite_content)
        print(f"[INFO] rsite content: {rsite_content}")

    def dock(self):
        if not (self._has_receptor and self._has_ligand):
            raise ValueError('Missing receptor or ligand.')
        self._prepare_lsite()

        cmd_hdock = [self.hdock_bin, "receptor.pdb", "ligand.pdb", "-lsite", "lsite.txt"]
        if self._epitope_sites is not None:
            self._prepare_rsite()
            cmd_hdock += ["-rsite", "rsite.txt"]
        subprocess.run(
            cmd_hdock,
            cwd=self.tmpdir.name, check=True
        )

        cmd_pl = [self.createpl_bin, "Hdock.out", "ligand_docked.pdb", "-lsite", "lsite.txt"]
        if self._epitope_sites is not None:
            self._prepare_rsite()
            cmd_pl += ["-rsite", "rsite.txt"]
        subprocess.run(
            cmd_pl, 
            cwd=self.tmpdir.name, check=True
        )
        return self._dump_complex_pdb()


if __name__ == '__main__':
    with HDockAntibody('hdock', 'createpl') as dock:
        dock.set_antigen('./data/dock/receptor.pdb', [DockSite('A', 991)])
        dock.set_antibody('./data/example_dock/3qhf_fv.pdb')
        print(dock.dock())

AbDesign/diffab/tools/dock/hdock.py
pengzhangzhi/ab_opt
import os
import re
import json
import shelve
from Bio import PDB
from typing import Optional, Tuple, List
from dataclasses import dataclass, field


@dataclass
class EvalTask:
    in_path: str
    ref_path: str
    info: dict
    structure: str
    name: str
    method: str
    cdr: str
    ab_chains: List

    residue_first: Optional[Tuple] = None
    residue_last: Optional[Tuple] = None
    
    scores: dict = field(default_factory=dict)

    def get_gen_biopython_model(self):
        parser = PDB.PDBParser(QUIET=True)
        return parser.get_structure(self.in_path, self.in_path)[0]

    def get_ref_biopython_model(self):
        parser = PDB.PDBParser(QUIET=True)
        return parser.get_structure(self.ref_path, self.ref_path)[0]

    def save_to_db(self, db: shelve.Shelf):
        db[self.in_path] = self

    def to_report_dict(self):
        return {
            'method': self.method,
            'structure': self.structure,
            'cdr': self.cdr,
            'filename': os.path.basename(self.in_path),
            **self.scores
        }


class TaskScanner:

    def __init__(self, root, postfix=None, db: Optional[shelve.Shelf]=None):
        super().__init__()
        self.root = root
        self.postfix = postfix
        self.visited = set()
        self.db = db
        if db is not None:
            for k in db.keys():
                self.visited.add(k)

    def _get_metadata(self, fpath):
        json_path = os.path.join(
            os.path.dirname(os.path.dirname(fpath)), 
            'metadata.json'
        )
        tag_name = os.path.basename(os.path.dirname(fpath))
        method_name = os.path.basename(
            os.path.dirname(os.path.dirname(os.path.dirname(fpath)))
        )
        try:
            antibody_chains = set()
            info = None
            with open(json_path, 'r') as f:
                metadata = json.load(f)
            for item in metadata['items']:
                if item['tag'] == tag_name:
                    info = item
                antibody_chains.add(item['residue_first'][0])
            if info is not None:
                info['antibody_chains'] = list(antibody_chains)
                info['structure'] = metadata['identifier']
                info['method'] = method_name
            return info
        except (json.JSONDecodeError, FileNotFoundError) as e:
            return None

    def scan(self) -> List[EvalTask]: 
        tasks = []
        if self.postfix is None or not self.postfix:
            input_fname_pattern = '^\d+\.pdb$'
            ref_fname = 'REF1.pdb'
        else:
            input_fname_pattern = f'^\d+\_{self.postfix}\.pdb$'
            ref_fname = f'REF1_{self.postfix}.pdb'
        for parent, _, files in os.walk(self.root):
            for fname in files:
                fpath = os.path.join(parent, fname)
                if not re.match(input_fname_pattern, fname):
                    continue
                if os.path.getsize(fpath) == 0:
                    continue
                if fpath in self.visited:
                    continue

                # Path to the reference structure
                ref_path = os.path.join(parent, ref_fname)
                if not os.path.exists(ref_path):
                    continue

                # CDR information
                info = self._get_metadata(fpath)
                if info is None:
                    continue
                tasks.append(EvalTask(
                    in_path = fpath,
                    ref_path = ref_path,
                    info = info,
                    structure = info['structure'],
                    name = info['name'],
                    method = info['method'],
                    cdr = info['tag'],
                    ab_chains = info['antibody_chains'],
                    residue_first = info.get('residue_first', None),
                    residue_last  = info.get('residue_last', None),
                ))
                self.visited.add(fpath)
        return tasks

AbDesign/diffab/tools/eval/base.py
pengzhangzhi/ab_opt
# pyright: reportMissingImports=false
import pyrosetta
from pyrosetta.rosetta.protocols.analysis import InterfaceAnalyzerMover
pyrosetta.init(' '.join([
    '-mute', 'all',
    '-use_input_sc',
    '-ignore_unrecognized_res',
    '-ignore_zero_occupancy', 'false',
    '-load_PDB_components', 'false',
    '-relax:default_repeats', '2',
    '-no_fconfig',
]))

from diffab.tools.eval.base import EvalTask


def pyrosetta_interface_energy(pdb_path, interface):
    pose = pyrosetta.pose_from_pdb(pdb_path)
    mover = InterfaceAnalyzerMover(interface)
    mover.set_pack_separated(True)
    mover.apply(pose)
    return pose.scores['dG_separated']


def eval_interface_energy(task: EvalTask):
    model_gen = task.get_gen_biopython_model()
    antigen_chains = set()
    for chain in model_gen:
        if chain.id not in task.ab_chains:
            antigen_chains.add(chain.id)
    antigen_chains = ''.join(list(antigen_chains))
    antibody_chains = ''.join(task.ab_chains)
    interface = f"{antibody_chains}_{antigen_chains}"

    dG_gen = pyrosetta_interface_energy(task.in_path, interface)
    dG_ref = pyrosetta_interface_energy(task.ref_path, interface)
    print(f"Energy: {dG_gen:.2f} vs {dG_ref:.2f} ({dG_gen - dG_ref:.2f})")
    task.scores.update({
        'dG_gen': dG_gen,
        'dG_ref': dG_ref,
        'ddG': dG_gen - dG_ref
    })
    return task

AbDesign/diffab/tools/eval/energy.py
pengzhangzhi/ab_opt
import os
import argparse
import ray
import shelve
import time
import pandas as pd
from typing import Mapping

from diffab.tools.eval.base import EvalTask, TaskScanner
from diffab.tools.eval.similarity import eval_similarity
from diffab.tools.eval.energy import eval_interface_energy


@ray.remote(num_cpus=1)
def evaluate(task, args):
    return evaluate_local(task, args.no_energy)


def evaluate_local(task, no_energy):
    funcs = []
    funcs.append(eval_similarity)
    if not no_energy:
        funcs.append(eval_interface_energy)
    for f in funcs:
        task = f(task)
    return task

def dump_db(db: Mapping[str, EvalTask], path):
    table = []
    for task in db.values():
        if 'abopt' in path and task.scores['seqid'] >= 100.0:
            # In abopt (Antibody Optimization) mode, ignore sequences identical to the wild-type
            continue
        table.append(task.to_report_dict())
    table = pd.DataFrame(table)
    table.to_csv(path, index=False, float_format='%.6f')
    return table

def run_energy(root, pfx='rosetta', no_energy=False):
    db_path = os.path.join(root, 'evaluation_db')
    with shelve.open(db_path) as db:
        scanner = TaskScanner(root=root, postfix=pfx, db=db)
        tasks = scanner.scan()
        for task in tasks:
            done_task = evaluate_local(task, no_energy)
            done_task.save_to_db(db)
        db.sync()
        dump_db(db, os.path.join(root, 'summary.csv'))
        
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, default='./results')
    parser.add_argument('--pfx', type=str, default='rosetta')
    parser.add_argument('--no_energy', action='store_true', default=False)
    args = parser.parse_args()
    ray.init()
    
    db_path = os.path.join(args.root, 'evaluation_db')
    with shelve.open(db_path) as db:
        scanner = TaskScanner(root=args.root, postfix=args.pfx, db=db)

        while True:        
            tasks = scanner.scan()
            futures = [evaluate.remote(t, args) for t in tasks]
            if len(futures) > 0:
                print(f'Submitted {len(futures)} tasks.')
            while len(futures) > 0:
                done_ids, futures = ray.wait(futures, num_returns=1)
                for done_id in done_ids:
                    done_task = ray.get(done_id)
                    done_task.save_to_db(db)
                    print(f'Remaining {len(futures)}. Finished {done_task.in_path}')
                db.sync()
            
            dump_db(db, os.path.join(args.root, 'summary.csv'))
            time.sleep(1.0)

if __name__ == '__main__':
    main()

AbDesign/diffab/tools/eval/run.py
pengzhangzhi/ab_opt
import numpy as np
from Bio.PDB import PDBParser, Selection
from Bio.PDB.Polypeptide import three_to_one
from Bio import pairwise2
from Bio.Align import substitution_matrices

from diffab.tools.eval.base import EvalTask


def reslist_rmsd(res_list1, res_list2):
    res_short, res_long = (res_list1, res_list2) if len(res_list1) < len(res_list2) else (res_list2, res_list1)
    M, N = len(res_short), len(res_long)

    def d(i, j):
        coord_i = np.array(res_short[i]['CA'].get_coord())
        coord_j = np.array(res_long[j]['CA'].get_coord())
        return ((coord_i - coord_j) ** 2).sum()

    SD = np.full([M, N], np.inf)
    for i in range(M):
        j = N - (M - i)
        SD[i, j] = sum([ d(i+k, j+k) for k in range(N-j) ])
    
    for j in range(N):
        SD[M-1, j] = d(M-1, j)

    for i in range(M-2, -1, -1):
        for j in range((N-(M-i))-1, -1, -1):
            SD[i, j] = min(
                d(i, j) + SD[i+1, j+1],
                SD[i, j+1]
            )

    min_SD = SD[0, :N-M+1].min()
    best_RMSD = np.sqrt(min_SD / M)
    return best_RMSD


def entity_to_seq(entity):
    seq = ''
    mapping = []
    for res in Selection.unfold_entities(entity, 'R'):
        try:
            seq += three_to_one(res.get_resname())
            mapping.append(res.get_id())
        except KeyError:
            pass
    assert len(seq) == len(mapping)
    return seq, mapping


def reslist_seqid(res_list1, res_list2):
    seq1, _ = entity_to_seq(res_list1)
    seq2, _ = entity_to_seq(res_list2)
    _, seq_id = align_sequences(seq1, seq2)
    return seq_id


def align_sequences(sequence_A, sequence_B, **kwargs):
    """
    Performs a global pairwise alignment between two sequences
    using the BLOSUM62 matrix and the Needleman-Wunsch algorithm
    as implemented in Biopython. Returns the alignment, the sequence
    identity and the residue mapping between both original sequences.
    """

    def _calculate_identity(sequenceA, sequenceB):
        """
        Returns the percentage of identical characters between two sequences.
        Assumes the sequences are aligned.
        """

        sa, sb, sl = sequenceA, sequenceB, len(sequenceA)
        matches = [sa[i] == sb[i] for i in range(sl)]
        seq_id = (100 * sum(matches)) / sl
        return seq_id

        # gapless_sl = sum([1 for i in range(sl) if (sa[i] != '-' and sb[i] != '-')])
        # gap_id = (100 * sum(matches)) / gapless_sl
        # return (seq_id, gap_id)

    #
    matrix = kwargs.get('matrix', substitution_matrices.load("BLOSUM62"))
    gap_open = kwargs.get('gap_open', -10.0)
    gap_extend = kwargs.get('gap_extend', -0.5)

    alns = pairwise2.align.globalds(sequence_A, sequence_B,
                                    matrix, gap_open, gap_extend,
                                    penalize_end_gaps=(False, False) )

    best_aln = alns[0]
    aligned_A, aligned_B, score, begin, end = best_aln

    # Calculate sequence identity
    seq_id = _calculate_identity(aligned_A, aligned_B)
    return (aligned_A, aligned_B), seq_id


def extract_reslist(model, residue_first, residue_last):
    assert residue_first[0] == residue_last[0]
    residue_first, residue_last = tuple(residue_first), tuple(residue_last)

    chain_id = residue_first[0]
    pos_first, pos_last = residue_first[1:], residue_last[1:]
    chain = model[chain_id]
    reslist = []
    for res in Selection.unfold_entities(chain, 'R'):
        pos_current = (res.id[1], res.id[2])
        if pos_first <= pos_current <= pos_last:
            reslist.append(res)
    return reslist


def eval_similarity(task: EvalTask):
    model_gen = task.get_gen_biopython_model()
    model_ref = task.get_ref_biopython_model()

    reslist_gen = extract_reslist(model_gen, task.residue_first, task.residue_last)
    reslist_ref = extract_reslist(model_ref, task.residue_first, task.residue_last)

    task.scores.update({
        'rmsd': reslist_rmsd(reslist_gen, reslist_ref),
        'seqid': reslist_seqid(reslist_gen, reslist_ref),
    })
    return task

AbDesign/diffab/tools/eval/similarity.py
pengzhangzhi/ab_opt
import os
import re
import json
from typing import Optional, Tuple, List
from dataclasses import dataclass


@dataclass
class RelaxTask:
    in_path: str
    current_path: str
    info: dict
    status: str

    flexible_residue_first: Optional[Tuple] = None
    flexible_residue_last: Optional[Tuple] = None

    def get_in_path_with_tag(self, tag):
        name, ext = os.path.splitext(self.in_path)
        new_path = f'{name}_{tag}{ext}'
        return new_path

    def set_current_path_tag(self, tag):
        new_path = self.get_in_path_with_tag(tag)
        self.current_path = new_path
        return new_path

    def check_current_path_exists(self):
        ok = os.path.exists(self.current_path)
        if not ok:
            self.mark_failure()
        if os.path.getsize(self.current_path) == 0:
            ok = False
            self.mark_failure()
            os.unlink(self.current_path)
        return ok

    def update_if_finished(self, tag):
        out_path = self.get_in_path_with_tag(tag)
        if os.path.exists(out_path) and os.path.getsize(out_path) > 0:
            print('Already finished', out_path)
            self.set_current_path_tag(tag)
            self.mark_success()
            return True
        return False

    def can_proceed(self):
        self.check_current_path_exists()
        return self.status != 'failed'

    def mark_success(self):
        self.status = 'success'

    def mark_failure(self):
        self.status = 'failed'

    def __str__(self):
        return f'{self.in_path} -> {self.current_path}'
    def __repr__(self):
        return str(self)
class TaskScanner:

    def __init__(self, root, final_postfix=None):
        super().__init__()
        self.root = root
        self.visited = set()
        self.final_postfix = final_postfix

    def _get_metadata(self, fpath):
        json_path = os.path.join(
            os.path.dirname(os.path.dirname(fpath)), 
            'metadata.json'
        )
        tag_name = os.path.basename(os.path.dirname(fpath))
        try:
            with open(json_path, 'r') as f:
                metadata = json.load(f)
            for item in metadata['items']:
                if item['tag'] == tag_name:
                    return item
        except (json.JSONDecodeError, FileNotFoundError) as e:
            return None
        return None

    def scan(self) -> List[RelaxTask]: 
        tasks = []
        input_fname_pattern = '^(?!.*patch)(^\d+\.pdb$|^REF\d\.pdb$)'
        for parent, _, files in os.walk(self.root):
            for fname in files:
                fpath = os.path.join(parent, fname)
                if not re.match(input_fname_pattern, fname):
                    continue
                if os.path.getsize(fpath) == 0:
                    continue
                if fpath in self.visited:
                    continue
                
                # If finished
                if self.final_postfix is not None:
                    fpath_name, fpath_ext = os.path.splitext(fpath)
                    fpath_final = f"{fpath_name}_{self.final_postfix}{fpath_ext}"
                    if os.path.exists(fpath_final):
                        continue

                # Get metadata
                info = self._get_metadata(fpath)
                if info is None:
                    continue
                    
                tasks.append(RelaxTask(
                    in_path = fpath,
                    current_path = fpath,
                    info = info,
                    status = 'created',
                    flexible_residue_first = info.get('residue_first', None),
                    flexible_residue_last  = info.get('residue_last', None),
                ))
                self.visited.add(fpath)
        return tasks

AbDesign/diffab/tools/relax/base.py
pengzhangzhi/ab_opt
import os
import time
import io
import logging
import pdbfixer
import openmm
from openmm import app as openmm_app
from openmm import unit
ENERGY = unit.kilocalories_per_mole
LENGTH = unit.angstroms

from diffab.tools.relax.base import RelaxTask


def current_milli_time():
    return round(time.time() * 1000)


def _is_in_the_range(ch_rs_ic, flexible_residue_first, flexible_residue_last):
    if ch_rs_ic[0] != flexible_residue_first[0]: return False
    r_first, r_last = tuple(flexible_residue_first[1:]), tuple(flexible_residue_last[1:])
    rs_ic = ch_rs_ic[1:]
    return r_first <= rs_ic <= r_last


class ForceFieldMinimizer(object):

    def __init__(self, stiffness=10.0, max_iterations=0, tolerance=2.39*unit.kilocalories_per_mole, platform='CUDA'):
        super().__init__()
        self.stiffness = stiffness
        self.max_iterations = max_iterations
        self.tolerance = tolerance
        assert platform in ('CUDA', 'CPU')
        self.platform = platform

    def _fix(self, pdb_str):
        fixer = pdbfixer.PDBFixer(pdbfile=io.StringIO(pdb_str))
        fixer.findNonstandardResidues()
        fixer.replaceNonstandardResidues()

        fixer.findMissingResidues()
        fixer.findMissingAtoms()
        fixer.addMissingAtoms(seed=0)
        fixer.addMissingHydrogens()

        out_handle = io.StringIO()
        openmm_app.PDBFile.writeFile(fixer.topology, fixer.positions, out_handle, keepIds=True)
        return out_handle.getvalue()

    def _get_pdb_string(self, topology, positions):
        with io.StringIO() as f:
            openmm_app.PDBFile.writeFile(topology, positions, f, keepIds=True)
            return f.getvalue()

    def _minimize(self, pdb_str, flexible_residue_first=None, flexible_residue_last=None):
        pdb = openmm_app.PDBFile(io.StringIO(pdb_str))

        force_field = openmm_app.ForceField("amber99sb.xml")
        constraints = openmm_app.HBonds
        system = force_field.createSystem(pdb.topology, constraints=constraints)

        # Add constraints to non-generated regions
        force = openmm.CustomExternalForce("0.5 * k * ((x-x0)^2 + (y-y0)^2 + (z-z0)^2)")
        force.addGlobalParameter("k", self.stiffness)
        for p in ["x0", "y0", "z0"]:
            force.addPerParticleParameter(p)
        
        if flexible_residue_first is not None and flexible_residue_last is not None:
            for i, a in enumerate(pdb.topology.atoms()):
                ch_rs_ic = (a.residue.chain.id, int(a.residue.id), a.residue.insertionCode)
                if not _is_in_the_range(ch_rs_ic, flexible_residue_first, flexible_residue_last) and a.element.name != "hydrogen":
                    force.addParticle(i, pdb.positions[i])
                
        system.addForce(force)

        # Set up the integrator and simulation
        integrator = openmm.LangevinIntegrator(0, 0.01, 0.0)
        platform = openmm.Platform.getPlatformByName("CUDA")
        simulation = openmm_app.Simulation(pdb.topology, system, integrator, platform)
        simulation.context.setPositions(pdb.positions)

        # Perform minimization
        ret = {}
        state = simulation.context.getState(getEnergy=True, getPositions=True)
        ret["einit"] = state.getPotentialEnergy().value_in_unit(ENERGY)
        ret["posinit"] = state.getPositions(asNumpy=True).value_in_unit(LENGTH)

        simulation.minimizeEnergy(maxIterations=self.max_iterations, tolerance=self.tolerance)

        state = simulation.context.getState(getEnergy=True, getPositions=True)
        ret["efinal"] = state.getPotentialEnergy().value_in_unit(ENERGY)
        ret["pos"] = state.getPositions(asNumpy=True).value_in_unit(LENGTH)
        ret["min_pdb"] = self._get_pdb_string(simulation.topology, state.getPositions())

        return ret['min_pdb'], ret

    def _add_energy_remarks(self, pdb_str, ret):
        pdb_lines = pdb_str.splitlines()
        pdb_lines.insert(1, "REMARK   1  FINAL ENERGY:   {:.3f} KCAL/MOL".format(ret['efinal']))
        pdb_lines.insert(1, "REMARK   1  INITIAL ENERGY: {:.3f} KCAL/MOL".format(ret['einit']))
        return "\n".join(pdb_lines)

    def __call__(self, pdb_str, flexible_residue_first=None, flexible_residue_last=None, return_info=True):
        if '\n' not in pdb_str and pdb_str.lower().endswith(".pdb"):
            with open(pdb_str) as f:
                pdb_str = f.read()

        pdb_fixed = self._fix(pdb_str)
        pdb_min, ret = self._minimize(pdb_fixed, flexible_residue_first, flexible_residue_last)
        pdb_min = self._add_energy_remarks(pdb_min, ret)
        if return_info:
            return pdb_min, ret
        else:
            return pdb_min


def run_openmm(task: RelaxTask):
    if not task.can_proceed() :
        return task
    if task.update_if_finished('openmm'):
        return task

    try:
        minimizer = ForceFieldMinimizer()
        with open(task.current_path, 'r') as f:
            pdb_str = f.read()

        pdb_min = minimizer(
            pdb_str = pdb_str,
            flexible_residue_first = task.flexible_residue_first,
            flexible_residue_last = task.flexible_residue_last,
            return_info = False,
        )
        out_path = task.set_current_path_tag('openmm')
        with open(out_path, 'w') as f:
            f.write(pdb_min)
        task.mark_success()
    except ValueError as e:
        logging.warning(
            f'{e.__class__.__name__}: {str(e)} ({task.current_path})'
        )        
        task.mark_failure()
    return task


AbDesign/diffab/tools/relax/openmm_relaxer.py
pengzhangzhi/ab_opt
# pyright: reportMissingImports=false
import os
import time
import pyrosetta
from pyrosetta.rosetta.protocols.relax import FastRelax
from pyrosetta.rosetta.core.pack.task import TaskFactory
from pyrosetta.rosetta.core.pack.task import operation
from pyrosetta.rosetta.core.select import residue_selector as selections
from pyrosetta.rosetta.core.select.movemap import MoveMapFactory, move_map_action
pyrosetta.init(' '.join([
    '-mute', 'all',
    '-use_input_sc',
    '-ignore_unrecognized_res',
    '-ignore_zero_occupancy', 'false',
    '-load_PDB_components', 'false',
    '-relax:default_repeats', '2',
    '-no_fconfig',
]))

from diffab.tools.relax.base import RelaxTask


def current_milli_time():
    return round(time.time() * 1000)


def parse_residue_position(p):
    icode = None
    if not p[-1].isnumeric():   # Has ICODE
        icode = p[-1]

    for i, c in enumerate(p):
        if c.isnumeric():
            break
    chain = p[:i]
    resseq = int(p[i:])

    if icode is not None:
        return chain, resseq, icode
    else:
        return chain, resseq


def get_scorefxn(scorefxn_name:str):
    """
    Gets the scorefxn with appropriate corrections.
    Taken from: https://gist.github.com/matteoferla/b33585f3aeab58b8424581279e032550
    """
    import pyrosetta

    corrections = {
        'beta_july15': False,
        'beta_nov16': False,
        'gen_potential': False,
        'restore_talaris_behavior': False,
    }
    if 'beta_july15' in scorefxn_name or 'beta_nov15' in scorefxn_name:
        # beta_july15 is ref2015
        corrections['beta_july15'] = True
    elif 'beta_nov16' in scorefxn_name:
        corrections['beta_nov16'] = True
    elif 'genpot' in scorefxn_name:
        corrections['gen_potential'] = True
        pyrosetta.rosetta.basic.options.set_boolean_option('corrections:beta_july15', True)
    elif 'talaris' in scorefxn_name:  #2013 and 2014
        corrections['restore_talaris_behavior'] = True
    else:
        pass
    for corr, value in corrections.items():
        pyrosetta.rosetta.basic.options.set_boolean_option(f'corrections:{corr}', value)
    return pyrosetta.create_score_function(scorefxn_name)


class RelaxRegion(object):
    
    def __init__(self, scorefxn='ref2015', max_iter=1000, subset='nbrs', move_bb=True):
        super().__init__()
        self.scorefxn = get_scorefxn(scorefxn)
        self.fast_relax = FastRelax()
        self.fast_relax.set_scorefxn(self.scorefxn)
        self.fast_relax.max_iter(max_iter)
        assert subset in ('all', 'target', 'nbrs')
        self.subset = subset
        self.move_bb = move_bb

    def __call__(self, pdb_path, flexible_residue_first, flexible_residue_last):
        pose = pyrosetta.pose_from_pdb(pdb_path)
        start_t = current_milli_time()
        original_pose = pose.clone()

        tf = TaskFactory()
        tf.push_back(operation.InitializeFromCommandline())
        tf.push_back(operation.RestrictToRepacking())   # Only allow residues to repack. No design at any position.

        # Create selector for the region to be relaxed
        # Turn off design and repacking on irrelevant positions
        if flexible_residue_first[-1] == ' ': 
            flexible_residue_first = flexible_residue_first[:-1]
        if flexible_residue_last[-1] == ' ':  
            flexible_residue_last  = flexible_residue_last[:-1]
        if self.subset != 'all':
            gen_selector = selections.ResidueIndexSelector()
            gen_selector.set_index_range(
                pose.pdb_info().pdb2pose(*flexible_residue_first), 
                pose.pdb_info().pdb2pose(*flexible_residue_last), 
            )
            nbr_selector = selections.NeighborhoodResidueSelector()
            nbr_selector.set_focus_selector(gen_selector)
            nbr_selector.set_include_focus_in_subset(True)

            if self.subset == 'nbrs':
                subset_selector = nbr_selector
            elif self.subset == 'target':
                subset_selector = gen_selector

            prevent_repacking_rlt = operation.PreventRepackingRLT()
            prevent_subset_repacking = operation.OperateOnResidueSubset(
                prevent_repacking_rlt, 
                subset_selector,
                flip_subset=True,
            )
            tf.push_back(prevent_subset_repacking)

        scorefxn = self.scorefxn
        fr = self.fast_relax

        pose = original_pose.clone()
        pos_list = pyrosetta.rosetta.utility.vector1_unsigned_long()
        for pos in range(pose.pdb_info().pdb2pose(*flexible_residue_first), pose.pdb_info().pdb2pose(*flexible_residue_last)+1):
            pos_list.append(pos)
        # basic_idealize(pose, pos_list, scorefxn, fast=True)

        mmf = MoveMapFactory()
        if self.move_bb: 
            mmf.add_bb_action(move_map_action.mm_enable, gen_selector)
        mmf.add_chi_action(move_map_action.mm_enable, subset_selector)
        mm  = mmf.create_movemap_from_pose(pose)

        fr.set_movemap(mm)
        fr.set_task_factory(tf)
        fr.apply(pose)

        e_before = scorefxn(original_pose)
        e_relax  = scorefxn(pose) 
        # print('\n\n[Finished in %.2f secs]' % ((current_milli_time() - start_t) / 1000))
        # print(' > Energy (before):    %.4f' % scorefxn(original_pose))
        # print(' > Energy (optimized): %.4f' % scorefxn(pose))
        return pose, e_before, e_relax


def run_pyrosetta(task: RelaxTask):
    if not task.can_proceed() :
        return task
    if task.update_if_finished('rosetta'):
        return task

    minimizer = RelaxRegion()
    pose_min, _, _ = minimizer(
        pdb_path = task.current_path,
        flexible_residue_first = task.flexible_residue_first,
        flexible_residue_last = task.flexible_residue_last,
    )

    out_path = task.set_current_path_tag('rosetta')
    pose_min.dump_pdb(out_path)
    task.mark_success()
    return task


def run_pyrosetta_fixbb(task: RelaxTask):
    if not task.can_proceed() :
        return task
    if task.update_if_finished('fixbb'):
        return task

    minimizer = RelaxRegion(move_bb=False)
    pose_min, _, _ = minimizer(
        pdb_path = task.current_path,
        flexible_residue_first = task.flexible_residue_first,
        flexible_residue_last = task.flexible_residue_last,
    )

    out_path = task.set_current_path_tag('fixbb')
    pose_min.dump_pdb(out_path)
    task.mark_success()
    return task

    


AbDesign/diffab/tools/relax/pyrosetta_relaxer.py
pengzhangzhi/ab_opt
import argparse
import ray
import time

from diffab.tools.relax.openmm_relaxer import run_openmm
from diffab.tools.relax.pyrosetta_relaxer import run_pyrosetta, run_pyrosetta_fixbb
from diffab.tools.relax.base import TaskScanner


@ray.remote(num_gpus=1/8, num_cpus=1)
def run_openmm_remote(task):
    return run_openmm(task)

def _run_openmm(task):
    return run_openmm(task)

@ray.remote(num_cpus=1)
def run_pyrosetta_remote(task):
    return run_pyrosetta(task)

def _run_pyrosetta(task):
    return run_pyrosetta(task)


@ray.remote(num_cpus=1)
def run_pyrosetta_fixbb_remote(task):
    return run_pyrosetta_fixbb(task)

def _run_pyrosetta_fixbb(task):
    return run_pyrosetta_fixbb(task)

@ray.remote
def pipeline_openmm_pyrosetta(task):
    funcs = [
        run_openmm_remote,
        run_pyrosetta_remote,
    ]
    for fn in funcs:
        task = fn.remote(task)
    return ray.get(task)

def _pipeline_openmm_pyrosetta(task):
    funcs = [
        _run_openmm,
        _run_pyrosetta,
    ]
    for fn in funcs:
        task = fn(task)
    return task

@ray.remote
def pipeline_pyrosetta(task):
    funcs = [
        run_pyrosetta_remote,
    ]
    for fn in funcs:
        task = fn.remote(task)
    return ray.get(task)

def _pipeline_pyrosetta(task):
    funcs = [
        _run_pyrosetta,
    ]
    for fn in funcs:
        task = fn(task)
    return task

@ray.remote
def pipeline_pyrosetta_fixbb(task):
    funcs = [
        run_pyrosetta_fixbb_remote,
    ]
    for fn in funcs:
        task = fn.remote(task)
    return ray.get(task)

def _pipeline_pyrosetta_fixbb(task):
    funcs = [
        _run_pyrosetta_fixbb,
    ]
    for fn in funcs:
        task = fn(task)
    return task


pipeline_dict = {
    'openmm_pyrosetta': pipeline_openmm_pyrosetta,
    'pyrosetta': pipeline_pyrosetta,
    'pyrosetta_fixbb': pipeline_pyrosetta_fixbb,
}

pipeline_dict_local = {
    'openmm_pyrosetta': _pipeline_openmm_pyrosetta,
    'pyrosetta': _pipeline_pyrosetta,
    'pyrosetta_fixbb': _pipeline_pyrosetta_fixbb,
}
def run_relax(root, pipeline_name='openmm_pyrosetta'):
    assert pipeline_name in pipeline_dict_local.keys()
    pipeline = pipeline_dict_local[pipeline_name]
    final_pfx = 'fixbb' if pipeline == pipeline_pyrosetta_fixbb else 'rosetta'
    scanner = TaskScanner(root, final_postfix=final_pfx)
    tasks = scanner.scan()
    for t in tasks:
        pipeline(t)
        

def main():
    ray.init()
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, default='./results')
    parser.add_argument('--pipeline', type=lambda s: pipeline_dict[s], default=pipeline_openmm_pyrosetta)
    args = parser.parse_args()

    final_pfx = 'fixbb' if args.pipeline == pipeline_pyrosetta_fixbb else 'rosetta'
    scanner = TaskScanner(args.root, final_postfix=final_pfx)
    while True:
        tasks = scanner.scan()
        futures = [args.pipeline.remote(t) for t in tasks]
        if len(futures) > 0:
            print(f'Submitted {len(futures)} tasks.')
        while len(futures) > 0:
            done_ids, futures = ray.wait(futures, num_returns=1)
            for done_id in done_ids:
                done_task = ray.get(done_id)
                print(f'Remaining {len(futures)}. Finished {done_task.current_path}')
        time.sleep(1.0)

if __name__ == '__main__':
    main()

AbDesign/diffab/tools/relax/run.py
pengzhangzhi/ab_opt
from .run import renumber

AbDesign/diffab/tools/renumber/__init__.py
pengzhangzhi/ab_opt
import argparse
import abnumber
from Bio import PDB
from Bio.PDB import Model, Chain, Residue, Selection
from Bio.Data import SCOPData
from typing import List, Tuple


def biopython_chain_to_sequence(chain: Chain.Chain):
    residue_list = Selection.unfold_entities(chain, 'R')
    seq = ''.join([SCOPData.protein_letters_3to1.get(r.resname, 'X') for r in residue_list])
    return seq, residue_list


def assign_number_to_sequence(seq):
    abchain = abnumber.Chain(seq, scheme='chothia')
    offset = seq.index(abchain.seq)
    if not (offset >= 0):
        raise ValueError(
            'The identified Fv sequence is not a subsequence of the original sequence.'
        )

    numbers = [None for _ in range(len(seq))]
    for i, (pos, aa) in enumerate(abchain):
        resseq = pos.number
        icode = pos.letter if pos.letter else ' '
        numbers[i+offset] = (resseq, icode)
    return numbers, abchain


def renumber_biopython_chain(chain_id, residue_list: List[Residue.Residue], numbers: List[Tuple[int, str]]):
    chain = Chain.Chain(chain_id)
    for residue, number in zip(residue_list, numbers):
        if number is None:
            continue
        residue = residue.copy()
        new_id = (residue.id[0], number[0], number[1])
        residue.id = new_id
        chain.add(residue)
    return chain


def renumber(in_pdb, out_pdb, return_other_chains=False):
    parser = PDB.PDBParser(QUIET=True)
    structure = parser.get_structure(None, in_pdb)
    model = structure[0]
    model_new = Model.Model(0)

    heavy_chains, light_chains, other_chains = [], [], []

    for chain in model:
        try:
            seq, reslist = biopython_chain_to_sequence(chain)
            numbers, abchain = assign_number_to_sequence(seq)
            chain_new = renumber_biopython_chain(chain.id, reslist, numbers)
            print(f'[INFO] Renumbered chain {chain_new.id} ({abchain.chain_type})')
            if abchain.chain_type == 'H':
                heavy_chains.append(chain_new.id)
            elif abchain.chain_type in ('K', 'L'):
                light_chains.append(chain_new.id)
        except abnumber.ChainParseError as e:
            print(f'[INFO] Chain {chain.id} does not contain valid Fv: {str(e)}')
            chain_new = chain.copy()
            other_chains.append(chain_new.id)
        model_new.add(chain_new)

    pdb_io = PDB.PDBIO()
    pdb_io.set_structure(model_new)
    pdb_io.save(out_pdb)
    if return_other_chains:
        return heavy_chains, light_chains, other_chains
    else:
        return heavy_chains, light_chains


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('in_pdb', type=str)
    parser.add_argument('out_pdb', type=str)
    args = parser.parse_args()

    renumber(args.in_pdb, args.out_pdb)

if __name__ == '__main__':
    main()

AbDesign/diffab/tools/renumber/run.py
pengzhangzhi/ab_opt
import os
import argparse
import copy
import json
from tqdm.auto import tqdm
from diffab.tools.runner.design_for_testset import ParseDict
from torch.utils.data import DataLoader

from diffab.datasets.custom import preprocess_antibody_structure
from diffab.models import get_model
from diffab.modules.common.geometry import reconstruct_backbone_partially
from diffab.modules.common.so3 import so3vec_to_rotation
from diffab.utils.inference import RemoveNative
from diffab.utils.protein.writers import save_pdb
from diffab.utils.train import recursive_to
from diffab.utils.misc import *
from diffab.utils.data import *
from diffab.utils.transforms import *
from diffab.utils.inference import *
from diffab.tools.renumber import renumber as renumber_antibody


def create_data_variants(config, structure_factory):
    structure = structure_factory()
    structure_id = structure['id']

    data_variants = []
    if config.mode == 'single_cdr':
        cdrs = sorted(list(set(find_cdrs(structure)).intersection(config.sampling.cdrs)))
        for cdr_name in cdrs:
            transform = Compose([
                MaskSingleCDR(cdr_name, augmentation=False),
                MergeChains(),
            ])
            data_var = transform(structure_factory())
            residue_first, residue_last = get_residue_first_last(data_var)
            data_variants.append({
                'data': data_var,
                'name': f'{structure_id}-{cdr_name}',
                'tag': f'{cdr_name}',
                'cdr': cdr_name,
                'residue_first': residue_first,
                'residue_last': residue_last,
            })
    elif config.mode == 'multiple_cdrs':
        cdrs = sorted(list(set(find_cdrs(structure)).intersection(config.sampling.cdrs)))
        transform = Compose([
            MaskMultipleCDRs(selection=cdrs, augmentation=False),
            MergeChains(),
        ])
        data_var = transform(structure_factory())
        data_variants.append({
            'data': data_var,
            'name': f'{structure_id}-MultipleCDRs',
            'tag': 'MultipleCDRs',
            'cdrs': cdrs,
            'residue_first': None,
            'residue_last': None,
        })
    elif config.mode == 'full':
        transform = Compose([
            MaskAntibody(),
            MergeChains(),
        ])
        data_var = transform(structure_factory())
        data_variants.append({
            'data': data_var,
            'name': f'{structure_id}-Full',
            'tag': 'Full',
            'residue_first': None,
            'residue_last': None,
        })
    elif config.mode == 'abopt':
        cdrs = sorted(list(set(find_cdrs(structure)).intersection(config.sampling.cdrs)))
        for cdr_name in cdrs:
            transform = Compose([
                MaskSingleCDR(cdr_name, augmentation=False),
                MergeChains(),
            ])
            data_var = transform(structure_factory())
            residue_first, residue_last = get_residue_first_last(data_var)
            for opt_step in config.sampling.optimize_steps:
                data_variants.append({
                    'data': data_var,
                    'name': f'{structure_id}-{cdr_name}-O{opt_step}',
                    'tag': f'{cdr_name}-O{opt_step}',
                    'cdr': cdr_name,
                    'opt_step': opt_step,
                    'residue_first': residue_first,
                    'residue_last': residue_last,
                })
    else:
        raise ValueError(f'Unknown mode: {config.mode}.')
    return data_variants


def design_for_pdb(args):
    # Load configs
    config, config_name = load_config(args.config)
    seed_all(args.seed if args.seed is not None else config.sampling.seed)

    # Structure loading
    data_id = os.path.basename(args.pdb_path)
    if args.no_renumber:
        pdb_path = args.pdb_path
    else:
        in_pdb_path = args.pdb_path
        out_pdb_path = os.path.splitext(in_pdb_path)[0] + '_chothia.pdb'
        heavy_chains, light_chains = renumber_antibody(in_pdb_path, out_pdb_path)
        pdb_path = out_pdb_path

        if args.heavy is None and len(heavy_chains) > 0:
            args.heavy = heavy_chains[0]
        if args.light is None and len(light_chains) > 0:
            args.light = light_chains[0]
    if args.heavy is None and args.light is None:
        raise ValueError("Neither heavy chain id (--heavy) or light chain id (--light) is specified.")
    get_structure = lambda: preprocess_antibody_structure({
        'id': data_id,
        'pdb_path': pdb_path,
        'heavy_id': args.heavy,
        # If the input is a nanobody, the light chain will be ignores
        'light_id': args.light,
    })

    # Logging
    structure_ = get_structure()
    structure_id = structure_['id']
    tag_postfix = '_%s' % args.tag if args.tag else ''
    log_dir = get_new_log_dir(
        os.path.join(args.out_root, config_name + tag_postfix), 
        prefix=data_id
    )
    logger = get_logger('sample', log_dir)
    logger.info(f'Data ID: {structure_["id"]}')
    logger.info(f'Results will be saved to {log_dir}')
    data_native = MergeChains()(structure_)
    save_pdb(data_native, os.path.join(log_dir, 'reference.pdb'))

    # Load checkpoint and model
    logger.info('Loading model config and checkpoints: %s' % (config.model.checkpoint))
    ckpt = torch.load(config.model.checkpoint, map_location='cpu')
    cfg_ckpt = ckpt['config']
    model = get_model(cfg_ckpt.model).to(args.device)
    lsd = model.load_state_dict(ckpt['model'])
    logger.info(str(lsd))

    # Make data variants
    data_variants = create_data_variants(
        config = config,
        structure_factory = get_structure,
    )

    # Save metadata
    metadata = {
        'identifier': structure_id,
        'index': data_id,
        'config': args.config,
        'items': [{kk: vv for kk, vv in var.items() if kk != 'data'} for var in data_variants],
    }
    with open(os.path.join(log_dir, 'metadata.json'), 'w') as f:
        json.dump(metadata, f, indent=2)

    # Start sampling
    collate_fn = CollateNodes(eight=False)
    inference_tfm = [ PatchAroundAnchor(), ]
    if 'abopt' not in config.mode:  # Don't remove native CDR in optimization mode
        inference_tfm.append(RemoveNative(
            remove_structure = config.sampling.sample_structure,
            remove_sequence = config.sampling.sample_sequence,
        ))
    inference_tfm = Compose(inference_tfm)

    for variant in data_variants:
        os.makedirs(os.path.join(log_dir, variant['tag']), exist_ok=True)
        logger.info(f"Start sampling for: {variant['tag']}")

        save_pdb(data_native, os.path.join(log_dir, variant['tag'], 'REF1.pdb'))       # w/  OpenMM minimization
    
        data_cropped = inference_tfm(
            copy.deepcopy(variant['data'])
        )
        data_list_repeat = [ data_cropped ] * config.sampling.num_samples
        loader = DataLoader(data_list_repeat, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn)
        
        count = 0
        for batch in tqdm(loader, desc=variant['name'], dynamic_ncols=True):
            torch.set_grad_enabled(False)
            model.eval()
            batch = recursive_to(batch, args.device)
            if 'abopt' in config.mode:
                # Antibody optimization starting from native
                traj_batch = model.optimize(batch, opt_step=variant['opt_step'], optimize_opt={
                    'pbar': True,
                    'sample_structure': config.sampling.sample_structure,
                    'sample_sequence': config.sampling.sample_sequence,
                })
            else:
                # De novo design
                traj_batch = model.sample(batch, sample_opt={
                    'pbar': True,
                    'sample_structure': config.sampling.sample_structure,
                    'sample_sequence': config.sampling.sample_sequence,
                    
                })

            aa_new = traj_batch[0][2]   # 0: Last sampling step. 2: Amino acid.
            pos_atom_new, mask_atom_new = reconstruct_backbone_partially(
                pos_ctx = batch['pos_heavyatom'],
                R_new = so3vec_to_rotation(traj_batch[0][0]),
                t_new = traj_batch[0][1],
                aa = aa_new,
                chain_nb = batch['chain_nb'],
                res_nb = batch['res_nb'],
                mask_atoms = batch['mask_heavyatom'],
                mask_recons = batch['generate_flag'],
            )
            aa_new = aa_new.cpu()
            pos_atom_new = pos_atom_new.cpu()
            mask_atom_new = mask_atom_new.cpu()

            for i in range(aa_new.size(0)):
                data_tmpl = variant['data']
                aa = apply_patch_to_tensor(data_tmpl['aa'], aa_new[i], data_cropped['patch_idx'])
                mask_ha = apply_patch_to_tensor(data_tmpl['mask_heavyatom'], mask_atom_new[i], data_cropped['patch_idx'])
                pos_ha  = (
                    apply_patch_to_tensor(
                        data_tmpl['pos_heavyatom'], 
                        pos_atom_new[i] + batch['origin'][i].view(1, 1, 3).cpu(), 
                        data_cropped['patch_idx']
                    )
                )

                save_path = os.path.join(log_dir, variant['tag'], '%04d.pdb' % (count, ))
                save_pdb({
                    'chain_nb': data_tmpl['chain_nb'],
                    'chain_id': data_tmpl['chain_id'],
                    'resseq': data_tmpl['resseq'],
                    'icode': data_tmpl['icode'],
                    # Generated
                    'aa': aa,
                    'mask_heavyatom': mask_ha,
                    'pos_heavyatom': pos_ha,
                }, path=save_path)
                # save_pdb({
                #     'chain_nb': data_cropped['chain_nb'],
                #     'chain_id': data_cropped['chain_id'],
                #     'resseq': data_cropped['resseq'],
                #     'icode': data_cropped['icode'],
                #     # Generated
                #     'aa': aa_new[i],
                #     'mask_heavyatom': mask_atom_new[i],
                #     'pos_heavyatom': pos_atom_new[i] + batch['origin'][i].view(1, 1, 3).cpu(),
                # }, path=os.path.join(log_dir, variant['tag'], '%04d_patch.pdb' % (count, )))
                count += 1

        logger.info('Finished.\n')


def args_from_cmdline():
    parser = argparse.ArgumentParser()
    parser.add_argument('pdb_path', type=str)
    parser.add_argument('--heavy', type=str, default=None, help='Chain id of the heavy chain.')
    parser.add_argument('--light', type=str, default=None, help='Chain id of the light chain.')
    parser.add_argument('--no_renumber', action='store_true', default=False)
    parser.add_argument('-c', '--config', type=str, default='./configs/test/codesign_single.yml')
    parser.add_argument('-o', '--out_root', type=str, default='./results')
    parser.add_argument('-t', '--tag', type=str, default='')
    parser.add_argument('-s', '--seed', type=int, default=None)
    parser.add_argument('-d', '--device', type=str, default='cuda')
    parser.add_argument('-b', '--batch_size', type=int, default=16)
    parser.add_argument(
        "--hydropathy_spec",
        metavar="KEY=VALUE",
        nargs="+",
        help="Set a number of key-value pairs "
        "(do not put spaces before or after the = sign). "
        "If a value contains spaces, you should define "
        "it with double quotes: "
        'foo="this is a sentence". Note that '
        "values are always treated as strings.",
        default=None,
        action=ParseDict,
    )
    parser.add_argument(
        "--charge_spec",
        metavar="KEY=VALUE",
        nargs="+",
        help="Set a number of key-value pairs "
        "(do not put spaces before or after the = sign). "
        "If a value contains spaces, you should define "
        "it with double quotes: "
        'foo="this is a sentence". Note that '
        "values are always treated as strings.",
        action=ParseDict,
        default=None,
    )
    args = parser.parse_args()
    return args


def args_factory(**kwargs):
    default_args = EasyDict(
        heavy = 'H',
        light = 'L',
        no_renumber = False,
        config = './configs/test/codesign_single.yml',
        out_root = './results',
        tag = '',
        seed = None,
        device = 'cuda',
        batch_size = 16
    )
    default_args.update(kwargs)
    return default_args


if __name__ == '__main__':
    design_for_pdb(args_from_cmdline())

AbDesign/diffab/tools/runner/design_for_pdb.py
pengzhangzhi/ab_opt
from collections import defaultdict
import os
import argparse
import copy
import json
from tqdm.auto import tqdm
from torch.utils.data import DataLoader

from diffab.datasets import get_dataset
from diffab.models import get_model
from diffab.modules.common.geometry import reconstruct_backbone_partially
from diffab.modules.common.so3 import so3vec_to_rotation
from diffab.tools.eval.run import run_energy
from diffab.tools.relax.run import run_relax
from diffab.utils.inference import RemoveNative
from diffab.utils.protein.constants import BBHeavyAtom
from diffab.utils.protein.writers import save_pdb
from diffab.utils.misc import *
from diffab.utils.data import *
from diffab.utils.train import eval_on_dataset,eval_sample
from diffab.utils.transforms import *
from diffab.utils.inference import *
from diffab.utils.val import create_data_variants, run_on_variant



def main():
    args = get_args()

    # Load configs
    config, config_name = load_config(args.config)
    seed_all(args.seed if args.seed is not None else config.sampling.seed)

    # Testset
    dataset = get_dataset(config.dataset.test)

    # Logging
    tag_postfix = '_%s' % args.tag if args.tag else ''
    # log_dir = get_new_log_dir(os.path.join(args.out_root, config_name + tag_postfix), prefix='%04d_%s' % (args.index, structure_['id']))
    logger = get_logger('sample')
    logger.info('Loading sampling configs: %s' % (args.config))
    ckpt_path = args.ckpt if args.ckpt is not None else config.model.checkpoint
    # Load checkpoint and model
    logger.info('Loading model config and checkpoints: %s' % (ckpt_path))
    ckpt = torch.load(ckpt_path, map_location='cpu')
    cfg_ckpt = ckpt['config']
    model = get_model(cfg_ckpt.model).to(args.device)
    lsd = model.load_state_dict(ckpt['model'])
    logger.info(str(lsd))

    save_dir = os.path.join(args.out_root, config_name + tag_postfix)
    if args.index is not None:
        get_structure = lambda: dataset[args.index]
        eval_sample(
            config, get_structure, model, logger, save_dir, 
            traj_idx=args.traj_idx, hydropathy_spec=args.hydropathy_spec, charge_spec=args.charge_spec
            )
    else:
        eval_on_dataset(config, dataset, model, logger, save_dir)
        run_relax(save_dir,pipeline_name='openmm_pyrosetta')
        run_energy(save_dir)
    
class ParseDict(argparse.Action):
    """copied from https://gist.github.com/fralau/061a4f6c13251367ef1d9a9a99fb3e8d?permalink_comment_id=4134590#gistcomment-4134590"""
    def __call__(self, parser, namespace, values, option_string=None):
        d = getattr(namespace, self.dest) or {}

        if values:
            for item in values:
                split_items = item.split("=", 1)
                key = split_items[
                    0
                ].strip()  # we remove blanks around keys, as is logical
                value = split_items[1]

                d[key] = value

        setattr(namespace, self.dest, d)
        
def get_args():
    parser = argparse.ArgumentParser()
    # parser.add_argument('index', type=int)
    parser.add_argument('-c', '--config', type=str, default='./configs/test/codesign_single.yml')
    parser.add_argument('-ck', '--ckpt', type=str, default='trained_models/ckpt.pt')
    parser.add_argument('-i', '--index', type=int, default=None)
    parser.add_argument('-o', '--out_root', type=str, default='test_results')
    parser.add_argument('-t', '--tag', type=str, default='')
    parser.add_argument('-s', '--seed', type=int, default=None)
    parser.add_argument('-d', '--device', type=str, default='cuda')
    parser.add_argument('-b', '--batch_size', type=int, default=1)
    parser.add_argument('-tr', '--traj_idx', type=int, default=0)
    parser.add_argument(
        "--hydropathy_spec",
        metavar="KEY=VALUE",
        nargs="+",
        help="Set a number of key-value pairs "
        "(do not put spaces before or after the = sign). "
        "If a value contains spaces, you should define "
        "it with double quotes: "
        'foo="this is a sentence". Note that '
        "values are always treated as strings.",
        default=None,
        action=ParseDict,
    )
    parser.add_argument(
        "--charge_spec",
        metavar="KEY=VALUE",
        nargs="+",
        help="Set a number of key-value pairs "
        "(do not put spaces before or after the = sign). "
        "If a value contains spaces, you should define "
        "it with double quotes: "
        'foo="this is a sentence". Note that '
        "values are always treated as strings.",
        action=ParseDict,
        default=None,
    )

    args = parser.parse_args()
    return args



if __name__ == '__main__':
    main()

AbDesign/diffab/tools/runner/design_for_testset.py
pengzhangzhi/ab_opt
import math
import torch
from torch.utils.data._utils.collate import default_collate
import torch.nn.functional as F
from torch_scatter import scatter_sum

DEFAULT_PAD_VALUES = {
    'aa': 21, 
    'chain_id': ' ', 
    'icode': ' ',
}

DEFAULT_NO_PADDING = {
    'origin',
}

class PaddingCollate(object):

    def __init__(self, length_ref_key='aa', pad_values=DEFAULT_PAD_VALUES, no_padding=DEFAULT_NO_PADDING, eight=True):
        super().__init__()
        self.length_ref_key = length_ref_key
        self.pad_values = pad_values
        self.no_padding = no_padding
        self.eight = eight

    @staticmethod
    def _pad_last(x, n, value=0):
        if isinstance(x, torch.Tensor):
            assert x.size(0) <= n
            if x.size(0) == n:
                return x
            pad_size = [n - x.size(0)] + list(x.shape[1:])
            pad = torch.full(pad_size, fill_value=value).to(x)
            return torch.cat([x, pad], dim=0)
        elif isinstance(x, list):
            pad = [value] * (n - len(x))
            return x + pad
        else:
            return x

    @staticmethod
    def _get_pad_mask(l, n):
        return torch.cat([
            torch.ones([l], dtype=torch.bool),
            torch.zeros([n-l], dtype=torch.bool)
        ], dim=0)

    @staticmethod
    def _get_common_keys(list_of_dict):
        keys = set(list_of_dict[0].keys())
        for d in list_of_dict[1:]:
            keys = keys.intersection(d.keys())
        return keys


    def _get_pad_value(self, key):
        if key not in self.pad_values:
            return 0
        return self.pad_values[key]

    def __call__(self, data_list):
        max_length = max([data[self.length_ref_key].size(0) for data in data_list])
        keys = self._get_common_keys(data_list)
        
        if self.eight:
            max_length = math.ceil(max_length / 8) * 8
        data_list_padded = []
        for data in data_list:
            data_padded = {
                k: self._pad_last(v, max_length, value=self._get_pad_value(k)) if k not in self.no_padding else v
                for k, v in data.items()
                if k in keys
            }
            data_padded['mask'] = self._get_pad_mask(data[self.length_ref_key].size(0), max_length)
            data_list_padded.append(data_padded)
        return default_collate(data_list_padded)


def apply_patch_to_tensor(x_full, x_patch, patch_idx):
    """
    Args:
        x_full:  (N, ...)
        x_patch: (M, ...)
        patch_idx:  (M, )
    Returns:
        (N, ...)
    """
    x_full = x_full.clone()
    x_full[patch_idx] = x_patch
    return x_full

class CollateNodes(object):
    def __init__(self,*args, **kwargs): 
        self.padding_collate = PaddingCollate(*args, **kwargs)
    def __call__(self, data_list):
        default_batch = self.padding_collate(data_list)
        coord_list = []
        aa_list = []
        batch_id_list = []
        fragment_id_list = []
        generate_flag_range_list = []
        for i,data in enumerate(data_list):
            seq_len = data['aa'].shape[0]
            coord = data['pos_heavyatom']
            fragment_type = data['fragment_type']
            batch_id = torch.full((seq_len,), i, dtype=torch.long)
            flag = torch.arange(seq_len)[data['generate_flag']]
            assert (flag[1:] - flag[:-1] == 1).all(), 'generate_flag must be continuous' 
            generate_flag_range_list.append(torch.tensor([flag[0], flag[-1]]))
            coord_list.append(coord)
            aa_list.append(data['aa'])
            batch_id_list.append(batch_id)  
            fragment_id_list.append(fragment_type)
        
        coord_tns = torch.cat(coord_list, dim=0) # [nb_all_residues, nb_atoms, 3]
        aa_tns = torch.cat(aa_list, dim=0) # [nb_all_residues]
        batch_id_tns = torch.cat(batch_id_list, dim=0)
        fragment_id_tns = torch.cat(fragment_id_list, dim=0)
        lengths = scatter_sum(torch.ones_like(batch_id_tns), batch_id_tns)  # [bs]
        offsets = F.pad(torch.cumsum(lengths, dim=0)[:-1], pad=(1, 0), value=0)  # [bs]
        # [bs, 2]
        generate_flag_range_tns = torch.stack(generate_flag_range_list, dim=0) + offsets.unsqueeze(-1)
        
        default_batch['graph'] = {
            'coord': coord_tns,
            'batch_id': batch_id_tns,
            'fragment_id': fragment_id_tns,
            "genearte_flag_range": generate_flag_range_tns,
            "aa": aa_tns,
        }
        return default_batch
AbDesign/diffab/utils/data.py
pengzhangzhi/ab_opt
import torch
from .protein import constants


def find_cdrs(structure):
    cdrs = []
    if structure['heavy'] is not None:
        flag = structure['heavy']['cdr_flag']
        if int(constants.CDR.H1) in flag:
            cdrs.append('H_CDR1')
        if int(constants.CDR.H2) in flag:
            cdrs.append('H_CDR2')
        if int(constants.CDR.H3) in flag:
            cdrs.append('H_CDR3')

    if structure['light'] is not None:
        flag = structure['light']['cdr_flag']
        if int(constants.CDR.L1) in flag:
            cdrs.append('L_CDR1')
        if int(constants.CDR.L2) in flag:
            cdrs.append('L_CDR2')
        if int(constants.CDR.L3) in flag:
            cdrs.append('L_CDR3')
    
    return cdrs


def get_residue_first_last(data):
    loop_flag = data['generate_flag']
    loop_idx = torch.arange(loop_flag.size(0))[loop_flag]
    idx_first, idx_last = loop_idx.min().item(), loop_idx.max().item()
    residue_first = (data['chain_id'][idx_first], data['resseq'][idx_first].item(), data['icode'][idx_first])
    residue_last = (data['chain_id'][idx_last], data['resseq'][idx_last].item(), data['icode'][idx_last])
    return residue_first, residue_last


class RemoveNative(object):

    def __init__(self, remove_structure, remove_sequence):
        super().__init__()
        self.remove_structure = remove_structure
        self.remove_sequence = remove_sequence

    def __call__(self, data):
        generate_flag = data['generate_flag'].clone()
        if self.remove_sequence:
            data['aa'] = torch.where(
                generate_flag, 
                torch.full_like(data['aa'], fill_value=int(constants.AA.UNK)),    # Is loop
                data['aa']
            )

        if self.remove_structure:
            data['pos_heavyatom'] = torch.where(
                generate_flag[:, None, None].expand(data['pos_heavyatom'].shape),
                torch.randn_like(data['pos_heavyatom']) * 10,
                data['pos_heavyatom']
            )

        return data
AbDesign/diffab/utils/inference.py
pengzhangzhi/ab_opt
from diffab.utils.data import apply_patch_to_tensor
from diffab.utils.protein.constants import BBHeavyAtom,ID2RES1,hydropathy2char,charge2char
import os
import time
import random
import logging
from typing import OrderedDict
import torch
import torch.linalg
import numpy as np
import yaml
from easydict import EasyDict
from glob import glob
import torch.nn.functional as F

from diffab.utils.protein.writers import save_pdb

class BlackHole(object):
    def __setattr__(self, name, value):
        pass

    def __call__(self, *args, **kwargs):
        return self

    def __getattr__(self, name):
        return self


class Counter(object):
    def __init__(self, start=0):
        super().__init__()
        self.now = start

    def step(self, delta=1):
        prev = self.now
        self.now += delta
        return prev


def get_logger(name, log_dir=None):
    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter('[%(asctime)s::%(name)s::%(levelname)s] %(message)s')

    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(logging.DEBUG)
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)

    if log_dir is not None:
        file_handler = logging.FileHandler(os.path.join(log_dir, 'log.txt'))
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

    return logger


def get_new_log_dir(root='./logs', prefix='', tag=''):
    fn = time.strftime('%Y_%m_%d__%H_%M_%S', time.localtime())
    if prefix != '':
        fn = prefix + '_' + fn
    if tag != '':
        fn = fn + '_' + tag
    log_dir = os.path.join(root, fn)
    os.makedirs(log_dir)
    return log_dir


def seed_all(seed):
    torch.backends.cudnn.deterministic = True
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)


def inf_iterator(iterable):
    iterator = iterable.__iter__()
    while True:
        try:
            yield iterator.__next__()
        except StopIteration:
            iterator = iterable.__iter__()


def log_hyperparams(writer, args):
    from torch.utils.tensorboard.summary import hparams
    vars_args = {k: v if isinstance(v, str) else repr(v) for k, v in vars(args).items()}
    exp, ssi, sei = hparams(vars_args, {})
    writer.file_writer.add_summary(exp)
    writer.file_writer.add_summary(ssi)
    writer.file_writer.add_summary(sei)


def int_tuple(argstr):
    return tuple(map(int, argstr.split(',')))


def str_tuple(argstr):
    return tuple(argstr.split(','))


def get_checkpoint_path(folder, it=None):
    if it is not None:
        return os.path.join(folder, '%d.pt' % it), it
    all_iters = list(map(lambda x: int(os.path.basename(x[:-3])), glob(os.path.join(folder, '*.pt'))))
    all_iters.sort()
    return os.path.join(folder, '%d.pt' % all_iters[-1]), all_iters[-1]


def load_config(config_path):
    with open(config_path, 'r') as f:
        config = EasyDict(yaml.safe_load(f))
    config_name = os.path.basename(config_path)[:os.path.basename(config_path).rfind('.')]
    return config, config_name


def extract_weights(weights: OrderedDict, prefix):
    extracted = OrderedDict()
    for k, v in weights.items():
        if k.startswith(prefix):
            extracted.update({
                k[len(prefix):]: v
            })
    return extracted


def current_milli_time():
    return round(time.time() * 1000)




def batchfy(nodes, lengths):
    """
    convert a sequence of samples into a structured batch.
    Args:
        nodes: (N, ...)
        lengths: (B, ), the length of each sample in the batch.
    Returns:
        batch: (B, n, ... ), n is the maximum length of the samples.
    """
    assert lengths.sum() == nodes.shape[0]
    n = lengths.max()
    b = len(lengths)
    # (B, n, ...)
    out = torch.zeros(
        (b, n, *nodes.shape[1:]), 
        dtype=nodes.dtype, device=nodes.device
        )
    offsets = F.pad(torch.cumsum(lengths, dim=0)[:-1], pad=(1, 0), value=0)
    for i, l in enumerate(lengths):
        offset = offsets[i]
        out[i, :l] = nodes[offset:offset+l]
    return out


def pair2edge(edge_idx, lengths, pair):
    """
    convert pair to edges.

    Args:
        edge_idx (torch.Tensor): [2, N]
        pair (torch.Tensor): [B, L, L, D]
        lengths (torch.Tensor): [B, ]
    Returns:
        edge (torch.Tensor): [N, D]
    """
    # [B+1, ] the offsets of each sample in the batch, pad zero at the beginning.
    B = len(lengths)
    offsets = F.pad(torch.cumsum(lengths, dim=0), pad=(1, 0), value=0)
    edge_feats = []
    for i in range(B):
        prev_offset = offsets[i]
        next_offset = offsets[i+1]
        idx = (edge_idx >= prev_offset) & (edge_idx < next_offset)
        local_edge_idx = edge_idx[idx].reshape(2,-1) - prev_offset
        edge_feat = pair[i][local_edge_idx[0],local_edge_idx[1]]
        edge_feats.append(edge_feat)
    edge_feats = torch.cat(edge_feats,dim=0)
    
    return edge_feats


def clash_loss(positions, seq_mask, chain_id, lit=3.0078, ep = 1e-6):
    """
    The clash loss penalizes too short distances only for residues from different chain.


    Args:
        positions (torch.Tensor): [B, N, 3]
        seq_mask (torch.Tensor): [B, N]
        chain_id (torch.Tensor): [B, N] the chain id of each residue
        lit (float): average shortest cdr to antigen distance. 
                    Defaults to 3.0078, which is calculated from the training set.
                    
        tolerance (float, optional): set to the standard deviation of the cdr2antigen distance.
                                    Defaults to 3.3694.
    Returns:
        loss (torch.Tensor): 
    """
    # [B, N, N]
    pair_different_chain_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2) 
    pair_mask = seq_mask.unsqueeze(1) * seq_mask.unsqueeze(2)
    mask = pair_mask.float() * pair_different_chain_mask.float()
    pair_dist = torch.cdist(positions, positions, p=2)
    pair_loss = F.relu(lit-pair_dist) * mask
    loss = pair_loss.sum() / ((pair_loss>0).sum() + ep) # average clash pairs
    # print(((pair_loss>0).sum(-1).sum(-1)/2).mean())
    # print(loss*10)
    return loss


def hotspot_distance_fn(x, hotspot_idx, cdr_idx,):
    """
    return the distance between cdr residues and the hotspots.
    

    Args:
        x (torch.Tensor): [B, N, 3]
        hotspot_idx (torch.Tensor): [B, X], X is 
        the number of hotspots
        cdr_idx (torch.Tensor): [B, Y], Y is the number of cdr residues
    Returns:
        pairwsie distance, [B, X, Y],
    """
    dist_mat = torch.cdist(x, x)
    cdr2hotspot_dist = dist_mat[:,cdr_idx, hotspot_idx]
    # cdr2hotspot_dist = dist_mat.gather(2, cdr_idx.unsqueeze(-1)).gather(1, hotspot_idx.unsqueeze(-1))
    return cdr2hotspot_dist
    

def save_each_sample(log_dir, variant, data_cropped, count, batch, aa_new, pos_atom_new, mask_atom_new, i):
    data_tmpl = variant['data']
    aa = apply_patch_to_tensor(data_tmpl['aa'], aa_new[i], data_cropped['patch_idx'])
    mask_ha = apply_patch_to_tensor(data_tmpl['mask_heavyatom'], mask_atom_new[i], data_cropped['patch_idx'])
    pos_ha  = (
                    apply_patch_to_tensor(
                        data_tmpl['pos_heavyatom'], 
                        pos_atom_new[i] + batch['origin'][i].view(1, 1, 3).cpu(), 
                        data_cropped['patch_idx']
                    )
                )

    save_path = os.path.join(log_dir, variant['tag'], '%04d.pdb' % (count, ))
    save_pdb({
        'chain_nb': data_tmpl['chain_nb'],
        'chain_id': data_tmpl['chain_id'],
        'resseq': data_tmpl['resseq'],
        'icode': data_tmpl['icode'],
        # Generated
        'aa': aa,
        'mask_heavyatom': mask_ha,
        'pos_heavyatom': pos_ha,
    }, path=save_path)
    generate_flags = data_tmpl["generate_flag"]
    native_atom_positions = data_tmpl["pos_heavyatom"][
                            ..., BBHeavyAtom.CA, :
                        ][generate_flags]
    pred_atom_positions = pos_ha[..., BBHeavyAtom.CA, :][generate_flags]
    rmsd = torch.sqrt(torch.mean(torch.sum(((native_atom_positions - pred_atom_positions) ** 2), dim=1)))
                
    native_aa = data_tmpl['aa'][generate_flags]
    pred_aa = aa[generate_flags]
    aa_recovery = torch.sum(native_aa == pred_aa) / native_aa.size(0)
    aa_seq = "".join([ID2RES1[int(aa_idx)] for aa_idx in pred_aa])
    native_aa_seq = "".join([ID2RES1[int(aa_idx)] for aa_idx in native_aa])
    hydropathy = "".join([hydropathy2char[int(label)] for label in data_cropped['hydropathy'][data_cropped['generate_flag']]])
    charge =  "".join([charge2char[int(label)] for label in data_cropped['charge'][data_cropped['generate_flag']]])
    # print(f"RMSD: {rmsd:.2f}, AA recovery: {aa_recovery:.2f}")
    save_pdb({
            'chain_nb': data_cropped['chain_nb'],
            'chain_id': data_cropped['chain_id'],
            'resseq': data_cropped['resseq'],
            'icode': data_cropped['icode'],
            # Generated
            'aa': aa_new[i],
            'mask_heavyatom': mask_atom_new[i],
            'pos_heavyatom': pos_atom_new[i] + batch['origin'][i].view(1, 1, 3).cpu(),
        }, 
        path=os.path.join(log_dir, variant['tag'], '%04d_patch.pdb' % (count, )))

    return {
        "rmsd": rmsd.item(),
        "aa_recovery": aa_recovery.item(),
        "aa_seq":aa_seq,
        "native_aa_seq": native_aa_seq,
        "hydropathy": hydropathy,
        "charge": charge,
    }, save_path
AbDesign/diffab/utils/misc.py
pengzhangzhi/ab_opt
import torch
import enum

class Hydropathy(enum.IntEnum):
    hydrophilic = 1
    moderate = 2
    unknown = 3
    
hydropathy2char = {
    Hydropathy.hydrophilic: '+',
    Hydropathy.moderate: '-',
    Hydropathy.unknown: '?',
}
char2hydropathy = {
    '+': Hydropathy.hydrophilic,
    '-': Hydropathy.moderate,
    '?': Hydropathy.unknown,
}

aa2hydropathy = {
    "C" : Hydropathy.moderate,
    "H": Hydropathy.moderate,
    "M": Hydropathy.moderate,
    
    "R": Hydropathy.hydrophilic,
    "K": Hydropathy.hydrophilic,
    "D": Hydropathy.hydrophilic,
    "E": Hydropathy.hydrophilic,
    "N": Hydropathy.hydrophilic,
    "Q": Hydropathy.hydrophilic,
    "S": Hydropathy.hydrophilic,
    "T": Hydropathy.hydrophilic,
    "Y": Hydropathy.hydrophilic,
    "W": Hydropathy.hydrophilic,
    "F": Hydropathy.hydrophilic,
    "L": Hydropathy.hydrophilic,
    "I": Hydropathy.hydrophilic,
    "V": Hydropathy.hydrophilic,
    "A": Hydropathy.hydrophilic,
    "G": Hydropathy.hydrophilic,
    "P": Hydropathy.hydrophilic,
}
class Charge(enum.IntEnum):
    positive = 1
    negtive = 2
    neutral = 3
    unknown = 4
    
charge2char = {
    Charge.positive: '+',
    Charge.negtive: '-',
    Charge.neutral: '=',
    Charge.unknown: '?',
}
char2charge = {
    '+': Charge.positive,
    '-': Charge.negtive,
    '=': Charge.neutral,
    '?': Charge.unknown,
}
aa2charge = {
    "R": Charge.positive,
    "K": Charge.positive,
    "H": Charge.positive,

    "D": Charge.negtive,
    "E": Charge.negtive,

    "S": Charge.neutral,
    "T": Charge.neutral,
    "N": Charge.neutral,
    "Q": Charge.neutral,
    "C": Charge.neutral,
    "G": Charge.neutral,
    "P": Charge.neutral,
    "A": Charge.neutral,
    "V": Charge.neutral,
    "I": Charge.neutral,
    "L": Charge.neutral,
    "M": Charge.neutral,
    "F": Charge.neutral,
    "Y": Charge.neutral,
    "W": Charge.neutral,
}
    
class CDR(enum.IntEnum):
    H1 = 1
    H2 = 2
    H3 = 3
    L1 = 4
    L2 = 5
    L3 = 6


class ChothiaCDRRange:
    H1 = (26, 32)
    H2 = (52, 56)
    H3 = (95, 102)
    
    L1 = (24, 34)
    L2 = (50, 56)
    L3 = (89, 97)

    @classmethod
    def to_cdr(cls, chain_type, resseq):
        assert chain_type in ('H', 'L')
        if chain_type == 'H':
            if cls.H1[0] <= resseq <= cls.H1[1]:
                return CDR.H1
            elif cls.H2[0] <= resseq <= cls.H2[1]:
                return CDR.H2
            elif cls.H3[0] <= resseq <= cls.H3[1]:
                return CDR.H3
        elif chain_type == 'L':
            if cls.L1[0] <= resseq <= cls.L1[1]:     # Chothia VH-CDR1
                return CDR.L1
            elif cls.L2[0] <= resseq <= cls.L2[1]:
                return CDR.L2
            elif cls.L3[0] <= resseq <= cls.L3[1]:
                return CDR.L3


class Fragment(enum.IntEnum):
    Heavy = 1
    Light = 2
    Antigen = 3

##
# Residue identities
"""
This is part of the OpenMM molecular simulation toolkit originating from
Simbios, the NIH National Center for Physics-Based Simulation of
Biological Structures at Stanford, funded under the NIH Roadmap for
Medical Research, grant U54 GM072970. See https://simtk.org.

Portions copyright (c) 2013 Stanford University and the Authors.
Authors: Peter Eastman
Contributors:

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
non_standard_residue_substitutions = {
    '2AS':'ASP', '3AH':'HIS', '5HP':'GLU', 'ACL':'ARG', 'AGM':'ARG', 'AIB':'ALA', 'ALM':'ALA', 'ALO':'THR', 'ALY':'LYS', 'ARM':'ARG',
    'ASA':'ASP', 'ASB':'ASP', 'ASK':'ASP', 'ASL':'ASP', 'ASQ':'ASP', 'AYA':'ALA', 'BCS':'CYS', 'BHD':'ASP', 'BMT':'THR', 'BNN':'ALA',
    'BUC':'CYS', 'BUG':'LEU', 'C5C':'CYS', 'C6C':'CYS', 'CAS':'CYS', 'CCS':'CYS', 'CEA':'CYS', 'CGU':'GLU', 'CHG':'ALA', 'CLE':'LEU', 'CME':'CYS',
    'CSD':'ALA', 'CSO':'CYS', 'CSP':'CYS', 'CSS':'CYS', 'CSW':'CYS', 'CSX':'CYS', 'CXM':'MET', 'CY1':'CYS', 'CY3':'CYS', 'CYG':'CYS',
    'CYM':'CYS', 'CYQ':'CYS', 'DAH':'PHE', 'DAL':'ALA', 'DAR':'ARG', 'DAS':'ASP', 'DCY':'CYS', 'DGL':'GLU', 'DGN':'GLN', 'DHA':'ALA',
    'DHI':'HIS', 'DIL':'ILE', 'DIV':'VAL', 'DLE':'LEU', 'DLY':'LYS', 'DNP':'ALA', 'DPN':'PHE', 'DPR':'PRO', 'DSN':'SER', 'DSP':'ASP',
    'DTH':'THR', 'DTR':'TRP', 'DTY':'TYR', 'DVA':'VAL', 'EFC':'CYS', 'FLA':'ALA', 'FME':'MET', 'GGL':'GLU', 'GL3':'GLY', 'GLZ':'GLY',
    'GMA':'GLU', 'GSC':'GLY', 'HAC':'ALA', 'HAR':'ARG', 'HIC':'HIS', 'HIP':'HIS', 'HMR':'ARG', 'HPQ':'PHE', 'HTR':'TRP', 'HYP':'PRO',
    'IAS':'ASP', 'IIL':'ILE', 'IYR':'TYR', 'KCX':'LYS', 'LLP':'LYS', 'LLY':'LYS', 'LTR':'TRP', 'LYM':'LYS', 'LYZ':'LYS', 'MAA':'ALA', 'MEN':'ASN',
    'MHS':'HIS', 'MIS':'SER', 'MLE':'LEU', 'MPQ':'GLY', 'MSA':'GLY', 'MSE':'MET', 'MVA':'VAL', 'NEM':'HIS', 'NEP':'HIS', 'NLE':'LEU',
    'NLN':'LEU', 'NLP':'LEU', 'NMC':'GLY', 'OAS':'SER', 'OCS':'CYS', 'OMT':'MET', 'PAQ':'TYR', 'PCA':'GLU', 'PEC':'CYS', 'PHI':'PHE',
    'PHL':'PHE', 'PR3':'CYS', 'PRR':'ALA', 'PTR':'TYR', 'PYX':'CYS', 'SAC':'SER', 'SAR':'GLY', 'SCH':'CYS', 'SCS':'CYS', 'SCY':'CYS',
    'SEL':'SER', 'SEP':'SER', 'SET':'SER', 'SHC':'CYS', 'SHR':'LYS', 'SMC':'CYS', 'SOC':'CYS', 'STY':'TYR', 'SVA':'SER', 'TIH':'ALA',
    'TPL':'TRP', 'TPO':'THR', 'TPQ':'ALA', 'TRG':'LYS', 'TRO':'TRP', 'TYB':'TYR', 'TYI':'TYR', 'TYQ':'TYR', 'TYS':'TYR', 'TYY':'TYR'
}


ressymb_to_resindex = {
    'A': 0, 'C': 1, 'D': 2, 'E': 3, 'F': 4,
    'G': 5, 'H': 6, 'I': 7, 'K': 8, 'L': 9,
    'M': 10, 'N': 11, 'P': 12, 'Q': 13, 'R': 14,
    'S': 15, 'T': 16, 'V': 17, 'W': 18, 'Y': 19,
    'X': 20,
}
ID2RES1 = {v: k for k, v in ressymb_to_resindex.items()}
aaidx2charge = {
    ressymb_to_resindex[symb] : charge for symb, charge in aa2charge.items()    
}
aaidx2hydropathy = {
    ressymb_to_resindex[symb] : hydrophobicity for symb, hydrophobicity in aa2hydropathy.items()
}

class AA(enum.IntEnum):
    ALA = 0; CYS = 1; ASP = 2; GLU = 3; PHE = 4
    GLY = 5; HIS = 6; ILE = 7; LYS = 8; LEU = 9
    MET = 10; ASN = 11; PRO = 12; GLN = 13; ARG = 14
    SER = 15; THR = 16; VAL = 17; TRP = 18; TYR = 19
    UNK = 20

    @classmethod
    def _missing_(cls, value):
        if isinstance(value, str) and len(value) == 3:      # three representation
            if value in non_standard_residue_substitutions:
                value = non_standard_residue_substitutions[value]
            if value in cls._member_names_:
                return getattr(cls, value)
        elif isinstance(value, str) and len(value) == 1:    # one representation
            if value in ressymb_to_resindex:
                return cls(ressymb_to_resindex[value])

        return super()._missing_(value)

    def __str__(self):
        return self.name

    @classmethod
    def is_aa(cls, value):
        return (value in ressymb_to_resindex) or \
            (value in non_standard_residue_substitutions) or \
            (value in cls._member_names_) or \
            (value in cls._member_map_.values())


num_aa_types = len(AA)

##
# Atom identities

class BBHeavyAtom(enum.IntEnum):
    N = 0; CA = 1; C = 2; O = 3; CB = 4; OXT=14;

max_num_heavyatoms = 15

# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
restype_to_heavyatom_names = {
    AA.ALA: ['N', 'CA', 'C', 'O', 'CB', '',    '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.ARG: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'NE',  'CZ',  'NH1', 'NH2', '',    '',    '', 'OXT'],
    AA.ASN: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'ND2', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.ASP: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'OD2', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.CYS: ['N', 'CA', 'C', 'O', 'CB', 'SG',  '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.GLN: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'NE2', '',    '',    '',    '',    '', 'OXT'],
    AA.GLU: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'OE2', '',    '',    '',    '',    '', 'OXT'],
    AA.GLY: ['N', 'CA', 'C', 'O', '',   '',    '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.HIS: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'ND1', 'CD2', 'CE1', 'NE2', '',    '',    '',    '', 'OXT'],
    AA.ILE: ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', 'CD1', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.LEU: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.LYS: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'CE',  'NZ',  '',    '',    '',    '',    '', 'OXT'],
    AA.MET: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'SD',  'CE',  '',    '',    '',    '',    '',    '', 'OXT'],
    AA.PHE: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  '',    '',    '', 'OXT'],
    AA.PRO: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.SER: ['N', 'CA', 'C', 'O', 'CB', 'OG',  '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.THR: ['N', 'CA', 'C', 'O', 'CB', 'OG1', 'CG2', '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.TRP: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'NE1', 'CE2', 'CE3', 'CZ2', 'CZ3', 'CH2', 'OXT'],
    AA.TYR: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  'OH',  '',    '', 'OXT'],
    AA.VAL: ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.UNK: ['',  '',   '',  '',  '',   '',    '',    '',    '',    '',    '',    '',    '',    '',    ''],
}
for names in restype_to_heavyatom_names.values(): assert len(names) == max_num_heavyatoms


backbone_atom_coordinates = {
    AA.ALA: [
        (-0.525, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, -0.0, -0.0),  # C
    ],
    AA.ARG: [
        (-0.524, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, -0.0),  # C
    ],
    AA.ASN: [
        (-0.536, 1.357, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, -0.0, -0.0),  # C
    ],
    AA.ASP: [
        (-0.525, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, 0.0, -0.0),  # C
    ],
    AA.CYS: [
        (-0.522, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.524, 0.0, 0.0),  # C
    ],
    AA.GLN: [
        (-0.526, 1.361, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, 0.0, 0.0),  # C
    ],
    AA.GLU: [
        (-0.528, 1.361, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, -0.0, -0.0),  # C
    ],
    AA.GLY: [
        (-0.572, 1.337, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.517, -0.0, -0.0),  # C
    ],
    AA.HIS: [
        (-0.527, 1.36, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, 0.0, 0.0),  # C
    ],
    AA.ILE: [
        (-0.493, 1.373, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, -0.0, -0.0),  # C
    ],
    AA.LEU: [
        (-0.52, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, -0.0),  # C
    ],
    AA.LYS: [
        (-0.526, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, 0.0, 0.0),  # C
    ],
    AA.MET: [
        (-0.521, 1.364, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, 0.0, 0.0),  # C
    ],
    AA.PHE: [
        (-0.518, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.524, 0.0, -0.0),  # C
    ],
    AA.PRO: [
        (-0.566, 1.351, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, -0.0, 0.0),  # C
    ],
    AA.SER: [
        (-0.529, 1.36, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, -0.0),  # C
    ],
    AA.THR: [
        (-0.517, 1.364, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, 0.0, -0.0),  # C
    ],
    AA.TRP: [
        (-0.521, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, 0.0),  # C
    ],
    AA.TYR: [
        (-0.522, 1.362, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.524, -0.0, -0.0),  # C
    ],
    AA.VAL: [
        (-0.494, 1.373, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, -0.0, -0.0),  # C
    ],
}

bb_oxygen_coordinate = {
    AA.ALA: (2.153, -1.062, 0.0),
    AA.ARG: (2.151, -1.062, 0.0),
    AA.ASN: (2.151, -1.062, 0.0),
    AA.ASP: (2.153, -1.062, 0.0),
    AA.CYS: (2.149, -1.062, 0.0),
    AA.GLN: (2.152, -1.062, 0.0),
    AA.GLU: (2.152, -1.062, 0.0),
    AA.GLY: (2.143, -1.062, 0.0),
    AA.HIS: (2.15, -1.063, 0.0),
    AA.ILE: (2.154, -1.062, 0.0),
    AA.LEU: (2.15, -1.063, 0.0),
    AA.LYS: (2.152, -1.062, 0.0),
    AA.MET: (2.15, -1.062, 0.0),
    AA.PHE: (2.15, -1.062, 0.0),
    AA.PRO: (2.148, -1.066, 0.0),
    AA.SER: (2.151, -1.062, 0.0),
    AA.THR: (2.152, -1.062, 0.0),
    AA.TRP: (2.152, -1.062, 0.0),
    AA.TYR: (2.151, -1.062, 0.0),
    AA.VAL: (2.154, -1.062, 0.0),
}

backbone_atom_coordinates_tensor = torch.zeros([21, 3, 3])
bb_oxygen_coordinate_tensor = torch.zeros([21, 3])

def make_coordinate_tensors():
    for restype, atom_coords in backbone_atom_coordinates.items():
        for atom_id, atom_coord in enumerate(atom_coords):
            backbone_atom_coordinates_tensor[restype][atom_id] = torch.FloatTensor(atom_coord)
    
    for restype, bb_oxy_coord in bb_oxygen_coordinate.items():
        bb_oxygen_coordinate_tensor[restype] = torch.FloatTensor(bb_oxy_coord)
make_coordinate_tensors()

AbDesign/diffab/utils/protein/constants.py
pengzhangzhi/ab_opt
import torch
from Bio.PDB import Selection
from Bio.PDB.Residue import Residue
from easydict import EasyDict

from .constants import (
    AA, max_num_heavyatoms,
    restype_to_heavyatom_names, 
    BBHeavyAtom
)


class ParsingException(Exception):
    pass


def _get_residue_heavyatom_info(res: Residue):
    pos_heavyatom = torch.zeros([max_num_heavyatoms, 3], dtype=torch.float)
    mask_heavyatom = torch.zeros([max_num_heavyatoms, ], dtype=torch.bool)
    restype = AA(res.get_resname())
    for idx, atom_name in enumerate(restype_to_heavyatom_names[restype]):
        if atom_name == '': continue
        if atom_name in res:
            pos_heavyatom[idx] = torch.tensor(res[atom_name].get_coord().tolist(), dtype=pos_heavyatom.dtype)
            mask_heavyatom[idx] = True
    return pos_heavyatom, mask_heavyatom


def parse_biopython_structure(entity, unknown_threshold=1.0, max_resseq=None):
    chains = Selection.unfold_entities(entity, 'C')
    chains.sort(key=lambda c: c.get_id())
    data = EasyDict({
        'chain_id': [],
        'resseq': [], 'icode': [], 'res_nb': [],
        'aa': [],
        'pos_heavyatom': [], 'mask_heavyatom': [],
    })
    tensor_types = {
        'resseq': torch.LongTensor,
        'res_nb': torch.LongTensor,
        'aa': torch.LongTensor,
        'pos_heavyatom': torch.stack,
        'mask_heavyatom': torch.stack,
    }

    count_aa, count_unk = 0, 0

    for i, chain in enumerate(chains):
        seq_this = 0   # Renumbering residues
        residues = Selection.unfold_entities(chain, 'R')
        residues.sort(key=lambda res: (res.get_id()[1], res.get_id()[2]))   # Sort residues by resseq-icode
        for _, res in enumerate(residues):
            resseq_this = int(res.get_id()[1])
            if max_resseq is not None and resseq_this > max_resseq:
                continue

            resname = res.get_resname()
            if not AA.is_aa(resname): continue
            if not (res.has_id('CA') and res.has_id('C') and res.has_id('N')): continue
            restype = AA(resname)
            count_aa += 1
            if restype == AA.UNK: 
                count_unk += 1
                continue

            # Chain info
            data.chain_id.append(chain.get_id())

            # Residue types
            data.aa.append(restype) # Will be automatically cast to torch.long

            # Heavy atoms
            pos_heavyatom, mask_heavyatom = _get_residue_heavyatom_info(res)
            data.pos_heavyatom.append(pos_heavyatom)
            data.mask_heavyatom.append(mask_heavyatom)

            # Sequential number
            resseq_this = int(res.get_id()[1])
            icode_this = res.get_id()[2]
            if seq_this == 0:
                seq_this = 1
            else:
                d_CA_CA = torch.linalg.norm(data.pos_heavyatom[-2][BBHeavyAtom.CA] - data.pos_heavyatom[-1][BBHeavyAtom.CA], ord=2).item()
                if d_CA_CA <= 4.0:
                    seq_this += 1
                else:
                    d_resseq = resseq_this - data.resseq[-1]
                    seq_this += max(2, d_resseq)

            data.resseq.append(resseq_this)
            data.icode.append(icode_this)
            data.res_nb.append(seq_this)

    if len(data.aa) == 0:
        raise ParsingException('No parsed residues.')

    if (count_unk / count_aa) >= unknown_threshold:
        raise ParsingException(
            f'Too many unknown residues, threshold {unknown_threshold:.2f}.'
        )

    seq_map = {}
    for i, (chain_id, resseq, icode) in enumerate(zip(data.chain_id, data.resseq, data.icode)):
        seq_map[(chain_id, resseq, icode)] = i

    for key, convert_fn in tensor_types.items():
        data[key] = convert_fn(data[key])

    return data, seq_map

AbDesign/diffab/utils/protein/parsers.py
pengzhangzhi/ab_opt
import torch
import warnings
from Bio import BiopythonWarning
from Bio.PDB import PDBIO
from Bio.PDB.StructureBuilder import StructureBuilder

from .constants import AA, restype_to_heavyatom_names


def save_pdb(data, path=None):
    """
    Args:
        data:   A dict that contains: `chain_nb`, `chain_id`, `aa`, `resseq`, `icode`,
                `pos_heavyatom`, `mask_heavyatom`.
    """

    def _mask_select(v, mask):
        if isinstance(v, str):
            return ''.join([s for i, s in enumerate(v) if mask[i]])
        elif isinstance(v, list):
            return [s for i, s in enumerate(v) if mask[i]]
        elif isinstance(v, torch.Tensor):
            return v[mask]
        else:
            return v

    def _build_chain(builder, aa_ch, pos_heavyatom_ch, mask_heavyatom_ch, chain_id_ch, resseq_ch, icode_ch):
        builder.init_chain(chain_id_ch[0])
        builder.init_seg('    ')

        for aa_res, pos_allatom_res, mask_allatom_res, resseq_res, icode_res in \
            zip(aa_ch, pos_heavyatom_ch, mask_heavyatom_ch, resseq_ch, icode_ch):
            if not AA.is_aa(aa_res.item()): 
                print('[Warning] Unknown amino acid type at %d%s: %r' % (resseq_res.item(), icode_res, aa_res.item()))
                continue
            restype = AA(aa_res.item())
            builder.init_residue(
                resname = str(restype),
                field = ' ',
                resseq = resseq_res.item(),
                icode = icode_res,
            )

            for i, atom_name in enumerate(restype_to_heavyatom_names[restype]):
                if atom_name == '': continue    # No expected atom
                if (~mask_allatom_res[i]).any(): continue     # Atom is missing
                if len(atom_name) == 1: fullname = ' %s  ' % atom_name
                elif len(atom_name) == 2: fullname = ' %s ' % atom_name
                elif len(atom_name) == 3: fullname = ' %s' % atom_name
                else: fullname = atom_name # len == 4
                builder.init_atom(atom_name, pos_allatom_res[i].tolist(), 0.0, 1.0, ' ', fullname,)

    warnings.simplefilter('ignore', BiopythonWarning)
    builder = StructureBuilder()
    builder.init_structure(0)
    builder.init_model(0)

    unique_chain_nb = data['chain_nb'].unique().tolist()
    for ch_nb in unique_chain_nb:
        mask = (data['chain_nb'] == ch_nb)
        aa = _mask_select(data['aa'], mask)
        pos_heavyatom = _mask_select(data['pos_heavyatom'], mask)
        mask_heavyatom = _mask_select(data['mask_heavyatom'], mask)
        chain_id = _mask_select(data['chain_id'], mask)
        resseq = _mask_select(data['resseq'], mask)
        icode = _mask_select(data['icode'], mask)

        _build_chain(builder, aa, pos_heavyatom, mask_heavyatom, chain_id, resseq, icode)
    
    structure = builder.get_structure()
    if path is not None:
        io = PDBIO()
        io.set_structure(structure)
        io.save(path)
    return structure

AbDesign/diffab/utils/protein/writers.py
pengzhangzhi/ab_opt
import pandas as pd
from collections import defaultdict
import json
import os
import numpy as np
import torch
import copy
import json
from tqdm.auto import tqdm
from torch.utils.data import DataLoader

from easydict import EasyDict
from diffab.modules.common.geometry import reconstruct_backbone_partially
from diffab.modules.common.so3 import so3vec_to_rotation
from diffab.utils.protein.writers import save_pdb
from diffab.utils.transforms import Compose

from diffab.utils.data import CollateNodes, PaddingCollate
from diffab.utils.inference import RemoveNative
from diffab.utils.transforms.label import Label
from diffab.utils.transforms.merge import MergeChains
from diffab.utils.transforms.patch import PatchAroundAnchor
from diffab.utils.val import create_data_variants,run_on_variant
from diffab.utils.protein.constants import char2hydropathy,char2charge
from .misc import BlackHole, save_each_sample


def get_optimizer(cfg, model):
    if cfg.type == 'adam':
        return torch.optim.Adam(
            model.parameters(),
            lr=cfg.lr,
            weight_decay=cfg.weight_decay,
            betas=(cfg.beta1, cfg.beta2, )
        )
    else:
        raise NotImplementedError('Optimizer not supported: %s' % cfg.type)


def get_scheduler(cfg, optimizer):
    if cfg.type is None:
        return BlackHole()
    elif cfg.type == 'plateau':
        return torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            factor=cfg.factor,
            patience=cfg.patience,
            min_lr=cfg.min_lr,
        )
    elif cfg.type == 'multistep':
        return torch.optim.lr_scheduler.MultiStepLR(
            optimizer,
            milestones=cfg.milestones,
            gamma=cfg.gamma,
        )
    elif cfg.type == 'exp':
        return torch.optim.lr_scheduler.ExponentialLR(
            optimizer,
            gamma=cfg.gamma,
        )
    elif cfg.type is None:
        return BlackHole()
    else:
        raise NotImplementedError('Scheduler not supported: %s' % cfg.type)


def get_warmup_sched(cfg, optimizer):
    if cfg is None: return BlackHole()
    lambdas = [lambda it : (it / cfg.max_iters) if it <= cfg.max_iters else 1 for _ in optimizer.param_groups]
    warmup_sched = torch.optim.lr_scheduler.LambdaLR(optimizer, lambdas)
    return warmup_sched


def log_losses(out, it, tag, logger=BlackHole(), writer=BlackHole(), others={}):
    logstr = '[%s] Iter %05d' % (tag, it)
    logstr += ' | loss %.4f' % out['overall'].item()
    for k, v in out.items():
        if k == 'overall': continue
        if isinstance(v, torch.Tensor): v = v.item()
        logstr += ' | loss(%s) %.4f' % (k, v)
    for k, v in others.items():
       logstr += ' | %s %2.4f' % (k, v)
    logger.info(logstr)

    for k, v in out.items():
        if k == 'overall':
            writer.add_scalar('%s/loss' % tag, v, it)
        else:
            writer.add_scalar('%s/loss_%s' % (tag, k), v, it)
    for k, v in others.items():
        writer.add_scalar('%s/%s' % (tag, k), v, it)
    writer.flush()


class ValidationLossTape(object):

    def __init__(self):
        super().__init__()
        self.accumulate = {}
        self.others = {}
        self.total = 0

    def update(self, out, n, others={}):
        self.total += n
        for k, v in out.items():
            if k not in self.accumulate:
                self.accumulate[k] = v.clone().detach()
            else:
                self.accumulate[k] += v.clone().detach()

        for k, v in others.items():
            if k not in self.others:
                self.others[k] = v.clone().detach()
            else:
                self.others[k] += v.clone().detach()
        

    def log(self, it, logger=BlackHole(), writer=BlackHole(), tag='val'):
        avg = EasyDict({k:v / self.total for k, v in self.accumulate.items()})
        avg_others = EasyDict({k:v / self.total for k, v in self.others.items()})
        log_losses(avg, it, tag, logger, writer, others=avg_others)
        return avg['overall']


def recursive_to(obj, device):
    if isinstance(obj, torch.Tensor):
        if device == 'cpu':
            return obj.cpu()
        try:
            return obj.cuda(device=device, non_blocking=True)
        except RuntimeError:
            return obj.to(device)
    elif isinstance(obj, list):
        return [recursive_to(o, device=device) for o in obj]
    elif isinstance(obj, tuple):
        return tuple(recursive_to(o, device=device) for o in obj)
    elif isinstance(obj, dict):
        return {k: recursive_to(v, device=device) for k, v in obj.items()}

    else:
        return obj


def reweight_loss_by_sequence_length(length, max_length, mode='sqrt'):
    if mode == 'sqrt':
        w = np.sqrt(length / max_length)
    elif mode == 'linear':
        w = length / max_length
    elif mode is None:
        w = 1.0
    else:
        raise ValueError('Unknown reweighting mode: %s' % mode)
    return w


def sum_weighted_losses(losses, weights):
    """
    Args:
        losses:     Dict of scalar tensors.
        weights:    Dict of weights.
    """
    if weights is None:
        # initialize weights to 1
        weights = {k:1.0 for k in losses.keys()}
    loss = 0
    for k in losses.keys():
        if not (isinstance(losses[k],torch.Tensor) and losses[k].requires_grad is True):
            # skip non-tensor or non-grad losses
            continue
        if weights.get(k) is None:
            weights[k] = 1
        loss = loss + weights[k] * losses[k]
    return loss


def count_parameters(model):
    return sum(p.numel() for p in model.parameters())


@torch.no_grad()
def eval_sample(config, get_structure, model, logger, save_dir, hydropathy_spec=None,charge_spec=None,  **kwargs):
    model.eval()
    sample_id = get_structure()['id']
    logger.info('Data ID: %s' % sample_id)

    sample_dir = os.path.join(save_dir, f"{sample_id}")
    os.makedirs(sample_dir, exist_ok=True)
    
    data_variants = create_data_variants(
        config = config,
        structure_factory = get_structure,
    )
    metadata = {
        'identifier': sample_id,
        'config': config,
        'items': [{kk: vv for kk, vv in var.items() if kk != 'data'} for var in data_variants],
    }
    with open(os.path.join(sample_dir, 'metadata.json'), 'w') as f:
        json.dump(metadata, f, indent=2)
    
    inference_tfm = [ 
                        PatchAroundAnchor(
                            config.antigen_size, 
                            config.initial_patch_size,
                            max_nb_hotspots=config.max_nb_hotspots,
                            is_train=False
                            ), 
                        Label(
                            hydropathy=config.hydropathy,
                            charge=config.charge,
                            drop_prob=config.drop_prob,
                        )
                    ]
    if 'abopt' not in config.mode:  # Don't remove native CDR in optimization mode
        inference_tfm.append(RemoveNative(
            remove_structure = config.sampling.sample_structure,
            remove_sequence = config.sampling.sample_sequence,
        ))
    inference_tfm = Compose(inference_tfm)
    
    data_native = MergeChains()(get_structure())
    config.sampling.num_samples = 1
    for variant in data_variants:
        os.makedirs(os.path.join(sample_dir, variant['tag']), exist_ok=True)
        logger.info(f"Start sampling for: {variant['tag']}")
        save_pdb(data_native, os.path.join(sample_dir, variant['tag'], 'REF1.pdb'))       # w/  OpenMM minimization
        
        data_cropped = inference_tfm(
                copy.deepcopy(variant['data'])
        )
        print('hydropathy_spec', hydropathy_spec)
        if hydropathy_spec is not None:
            hydropathy_labels = data_cropped['hydropathy'][data_cropped['generate_flag']]
            print("before:", hydropathy_labels)
            for idx, label in hydropathy_spec.items():
                hydropathy_labels[int(idx)] = char2hydropathy[str(label)]
            data_cropped['hydropathy'][data_cropped['generate_flag']] = hydropathy_labels
            print("after:", data_cropped['hydropathy'][data_cropped['generate_flag']])
        if charge_spec is not None:
            charge_labels = data_cropped['charge'][data_cropped['generate_flag']]
            for idx, label in charge_spec.items():
                charge_labels[int(idx)] = char2charge[str(label)]
            data_cropped['charge'][data_cropped['generate_flag']] = charge_labels
        batch_size = min(100, config.sampling.num_samples)
        assert config.sampling.num_samples == 1, 'sample once and save all the trajectories'
        data_list_repeat = [ data_cropped ] * config.sampling.num_samples
        collate_fn = CollateNodes(eight=False)
        loader = DataLoader(data_list_repeat, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
        
        for batch in tqdm(loader, desc=variant['name'], dynamic_ncols=True):
            model.eval()
            device = next(model.parameters()).device
            batch = recursive_to(batch, device)
            if 'abopt' in config.mode:
                # Antibody optimization starting from native
                traj_batch = model.optimize(batch, opt_step=variant['opt_step'], optimize_opt={
                            'pbar': True,
                            'sample_structure': config.sampling.sample_structure,
                            'sample_sequence': config.sampling.sample_sequence,
                        })
            else:
                        # De novo design
                traj_batch = model.sample(batch,
                        sample_opt={
                            'pbar': True,
                            'sample_structure': config.sampling.sample_structure,
                            'sample_sequence': config.sampling.sample_sequence,
                        },
                        
                        )
            traj_pdb_str = ''
            for traj_idx in (traj_batch):
                r_new = so3vec_to_rotation(traj_batch[traj_idx][0]).to(device)
                t_new = (traj_batch[traj_idx][1]).to(device)
                aa_new = (traj_batch[traj_idx][2]).to(device)   # 2: Amino acid.
                pos_atom_new, mask_atom_new = reconstruct_backbone_partially(
                            pos_ctx = batch['pos_heavyatom'],
                            R_new = r_new,
                            t_new = t_new,
                            aa = aa_new, 
                            chain_nb = batch['chain_nb'],
                            res_nb = batch['res_nb'],
                            mask_atoms = batch['mask_heavyatom'],
                            mask_recons = batch['generate_flag'],
                        )
                aa_new = aa_new.cpu()
                pos_atom_new = pos_atom_new.cpu()
                mask_atom_new = mask_atom_new.cpu()
                result,pdb_path = save_each_sample(sample_dir, variant, data_cropped, traj_idx, batch, aa_new, pos_atom_new, mask_atom_new, i=0)
                result = {variant['tag']+"-"+k:torch.tensor(v).mean().item() if isinstance(v, (torch.Tensor, float)) else v for k, v in result.items() }
                with open(pdb_path, 'r') as f:
                    traj_pdb_str += f'MODEL     {traj_idx+1}'+f.read()+'ENDMDL\n'
                    
            traj_path = os.path.join(sample_dir, f"traj.pdb")
            with open(traj_path, 'w') as f:
                f.write(traj_pdb_str)
                print(f"Saved trajectory to {traj_path}")
            logger.info(f"traj: {traj_idx} "+", ".join(f"{k}: {v:.2f}" if isinstance(v,float) else f"{k}: {v}" for k, v in result.items()))
        

@torch.no_grad()
def eval_on_dataset(config, dataset, model, logger, save_dir):
    """ evaluate on val or test set. """
    datset_size = len(dataset)
    model.eval()
    dataset_results = defaultdict(list)
    id_lst = []
    for i in range(datset_size):
        # evaluate each sample
        get_structure = lambda: dataset[i]
        sample_id = get_structure()['id']
        logger.info('Data ID: %s' % sample_id)
    
        sample_dir = os.path.join(save_dir, f"{str(i)}-{sample_id}")
        os.makedirs(sample_dir, exist_ok=True)
        
        data_variants = create_data_variants(
            config = config,
            structure_factory = get_structure,
        )
        metadata = {
            'identifier': sample_id,
            'index': i,
            'config': config,
            'items': [{kk: vv for kk, vv in var.items() if kk != 'data'} for var in data_variants],
        }
        with open(os.path.join(sample_dir, 'metadata.json'), 'w') as f:
            json.dump(metadata, f, indent=2)
        
        inference_tfm = [ 
                         PatchAroundAnchor(
                             config.antigen_size, 
                             config.initial_patch_size,
                             max_nb_hotspots=config.max_nb_hotspots,
                             is_train=False
                             ), 
                        Label(
                            hydropathy=config.hydropathy, charge=config.charge,drop_prob=config.drop_prob,
                        )
                        ]
        if 'abopt' not in config.mode:  # Don't remove native CDR in optimization mode
            inference_tfm.append(RemoveNative(
                remove_structure = config.sampling.sample_structure,
                remove_sequence = config.sampling.sample_sequence,
            ))
        inference_tfm = Compose(inference_tfm)
        
        data_native = MergeChains()(get_structure())
        for variant in data_variants:
            result = run_on_variant(config, sample_dir, logger, data_native, model, inference_tfm, variant)
            
            result = {variant['tag']+"-"+k:torch.tensor(v).mean().item() if isinstance(v[0],(float,torch.Tensor)) else v[-1] for k, v in result.items()}
            
            logger.info(f"idx: {i}, "+", ".join(f"{k}: {v:.2f}" if isinstance(v,float) else f"{k}: {v}" for k, v in result.items()))
            for k,v in result.items():
                dataset_results[k].append(v)
                
        if len(data_variants) != 0: id_lst.append(sample_id)
        
    df = pd.DataFrame(dataset_results, index=id_lst)
    # df.loc['mean'] = df.mean()
    df.to_csv(os.path.join(save_dir, 'results.csv'), float_format='%.2f')
    print(df)
    return dict(df.mean())
AbDesign/diffab/utils/train.py
pengzhangzhi/ab_opt
# Transforms
from .mask import MaskSingleCDR, MaskMultipleCDRs, MaskAntibody
from .merge import MergeChains
from .patch import PatchAroundAnchor
from .filter_structure import FilterStructure
from .label import Label

# Factory
from ._base import get_transform, Compose

AbDesign/diffab/utils/transforms/__init__.py
pengzhangzhi/ab_opt
import copy
import torch
# from torchvision.transforms import Compose


_TRANSFORM_DICT = {}

class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms
    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img
    def __repr__(self) -> str:
        format_string = self.__class__.__name__ + "("
        for t in self.transforms:
            format_string += "\n"
            format_string += f"    {t}"
        format_string += "\n)"
        return format_string

    
def register_transform(name):
    def decorator(cls):
        _TRANSFORM_DICT[name] = cls
        return cls
    return decorator


def get_transform(cfg):
    if cfg is None or len(cfg) == 0:
        return None
    tfms = []
    for t_dict in cfg:
        t_dict = copy.deepcopy(t_dict)
        cls = _TRANSFORM_DICT[t_dict.pop('type')]
        tfms.append(cls(**t_dict))
    return Compose(tfms)


def _index_select(v, index, n):
    if isinstance(v, torch.Tensor) and v.size(0) == n:
        return v[index]
    elif isinstance(v, list) and len(v) == n:
        return [v[i] for i in index]
    else:
        return v


def _index_select_data(data, index):
    return {
        k: _index_select(v, index, data['aa'].size(0))
        for k, v in data.items()
    }


def _mask_select(v, mask):
    if isinstance(v, torch.Tensor) and v.size(0) == mask.size(0):
        return v[mask]
    elif isinstance(v, list) and len(v) == mask.size(0):
        return [v[i] for i, b in enumerate(mask) if b]
    else:
        return v


def _mask_select_data(data, mask):
    return {
        k: _mask_select(v, mask)
        for k, v in data.items()
    }

AbDesign/diffab/utils/transforms/_base.py
pengzhangzhi/ab_opt
import torch

from ..protein import constants
from ._base import register_transform

class MissingChainException(Exception):
    pass

@register_transform('filter_structure')
class FilterStructure(object):

    def __init__(self, must_have_heavy=False, must_have_light=False, must_have_antigen=False):
        super().__init__()
        self.must_have_heavy = must_have_heavy
        self.must_have_light = must_have_light
        self.must_have_antigen = must_have_antigen
        
    

    def __call__(self, structure):
        if self.must_have_heavy and structure['heavy'] is None:
            raise MissingChainException('heavy chain missing')
        if self.must_have_light and structure['light'] is None:
            raise MissingChainException('light chain missing')
        if self.must_have_antigen and structure['antigen'] is None:
            raise MissingChainException('antigen missing')
        
        return structure


AbDesign/diffab/utils/transforms/filter_structure.py
pengzhangzhi/ab_opt
import torch

from ._base import _mask_select_data, register_transform
from ..protein.constants import Hydropathy, Charge, aaidx2charge,aaidx2hydropathy


@register_transform('aa_label')
class Label(object):

    def __init__(self, hydropathy=False, charge=False,drop_prob=0.0):
        super().__init__()
        self.hydropathy = hydropathy
        self.charge = charge
        self.drop_prob = drop_prob
    def __call__(self, data):
        aa_flags = data['generate_flag']
        aa = data['aa']
        
        hydropathy_flags = torch.zeros_like(aa).fill_(Hydropathy.unknown)
        if self.hydropathy:
            hydropathy_label = torch.tensor([aaidx2hydropathy[int(aaidx)] for aaidx in aa[aa_flags]],dtype=torch.int64)
            # randomly drop some labels with probability drop_prob
            if self.drop_prob > 0.0:
                drop_mask = torch.rand_like(hydropathy_label.float()) < self.drop_prob 
                hydropathy_label[drop_mask] = Hydropathy.unknown
            hydropathy_flags[aa_flags] = hydropathy_label
        data['hydropathy'] = hydropathy_flags
        
        charge_flags = torch.zeros_like(aa).fill_(Charge.unknown)
        if self.charge:
            charge_label = torch.tensor([aaidx2charge[int(aaidx)] for aaidx in aa[aa_flags]], dtype=torch.int64)
            # randomly drop some labels with probability drop_prob
            if self.drop_prob > 0.0:
                drop_mask = torch.rand_like(charge_label.float()) < self.drop_prob
                charge_label[drop_mask] = Charge.unknown
            charge_flags[aa_flags] = charge_label
        data['charge'] = charge_flags
        return data
AbDesign/diffab/utils/transforms/label.py
pengzhangzhi/ab_opt
import torch
import random
from typing import List, Optional

from ..protein import constants
from ._base import register_transform
from diffab.utils.transforms.filter_structure import MissingChainException


def random_shrink_extend(flag, min_length=5, shrink_limit=1, extend_limit=2):
    first, last = continuous_flag_to_range(flag)
    length = flag.sum().item()
    if (length - 2*shrink_limit) < min_length:
        shrink_limit = 0
    first_ext = max(0, first-random.randint(-shrink_limit, extend_limit))
    last_ext = min(last+random.randint(-shrink_limit, extend_limit), flag.size(0)-1)
    flag_ext = flag.clone()
    flag_ext[first_ext : last_ext+1] = True
    return flag_ext


def continuous_flag_to_range(flag):
    first = (torch.arange(0, flag.size(0))[flag]).min().item()
    last = (torch.arange(0, flag.size(0))[flag]).max().item()
    return first, last


@register_transform('mask_single_cdr')
class MaskSingleCDR(object):

    def __init__(self, selection=None, augmentation=True):
        super().__init__()
        cdr_str_to_enum = {
            'H1': constants.CDR.H1,
            'H2': constants.CDR.H2,
            'H3': constants.CDR.H3,
            'L1': constants.CDR.L1,
            'L2': constants.CDR.L2,
            'L3': constants.CDR.L3,
            'H_CDR1': constants.CDR.H1,
            'H_CDR2': constants.CDR.H2,
            'H_CDR3': constants.CDR.H3,
            'L_CDR1': constants.CDR.L1,
            'L_CDR2': constants.CDR.L2,
            'L_CDR3': constants.CDR.L3,
            'CDR3': 'CDR3',     # H3 first, then fallback to L3
        }
        assert selection is None or selection in cdr_str_to_enum
        self.selection = cdr_str_to_enum.get(selection, None)
        self.augmentation = augmentation

    def perform_masking_(self, data, selection=None):
        cdr_flag = data['cdr_flag']

        if selection is None:
            cdr_all = cdr_flag[cdr_flag > 0].unique().tolist()
            cdr_to_mask = random.choice(cdr_all)
        else:
            cdr_to_mask = selection

        cdr_to_mask_flag = (cdr_flag == cdr_to_mask)
        if self.augmentation:
            cdr_to_mask_flag = random_shrink_extend(cdr_to_mask_flag)

        cdr_first, cdr_last = continuous_flag_to_range(cdr_to_mask_flag)
        left_idx = max(0, cdr_first-1)
        right_idx = min(data['aa'].size(0)-1, cdr_last+1)
        anchor_flag = torch.zeros(data['aa'].shape, dtype=torch.bool)
        anchor_flag[left_idx] = True
        anchor_flag[right_idx] = True

        data['generate_flag'] = cdr_to_mask_flag
        data['anchor_flag'] = anchor_flag

    def __call__(self, structure):
        if self.selection is None:
            ab_data = []
            if structure['heavy'] is not None:
                ab_data.append(structure['heavy'])
            if structure['light'] is not None:
                ab_data.append(structure['light'])
            data_to_mask = random.choice(ab_data)
            sel = None
        elif self.selection in (constants.CDR.H1, constants.CDR.H2, constants.CDR.H3, ) and structure['heavy'] is not None:
            data_to_mask = structure['heavy']
            sel = int(self.selection)
        elif self.selection in (constants.CDR.L1, constants.CDR.L2, constants.CDR.L3, ) and structure['light'] is not None:
            data_to_mask = structure['light']
            sel = int(self.selection)
        elif self.selection == 'CDR3':
            if structure['heavy'] is not None:
                data_to_mask = structure['heavy']
                sel = constants.CDR.H3
            else:
                data_to_mask = structure['light']
                sel = constants.CDR.L3
        else:
            raise MissingChainException(f"select {self.selection}, but this chain is missing")
        self.perform_masking_(data_to_mask, selection=sel)
        return structure


@register_transform('mask_multiple_cdrs')
class MaskMultipleCDRs(object):
    
    def __init__(self, selection: Optional[List[str]]=None, augmentation=True):
        super().__init__()
        cdr_str_to_enum = {
            'H1': constants.CDR.H1,
            'H2': constants.CDR.H2,
            'H3': constants.CDR.H3,
            'L1': constants.CDR.L1,
            'L2': constants.CDR.L2,
            'L3': constants.CDR.L3,
            'H_CDR1': constants.CDR.H1,
            'H_CDR2': constants.CDR.H2,
            'H_CDR3': constants.CDR.H3,
            'L_CDR1': constants.CDR.L1,
            'L_CDR2': constants.CDR.L2,
            'L_CDR3': constants.CDR.L3,
        }
        if selection is not None:
            self.selection = [cdr_str_to_enum[s] for s in selection]
        else:
            self.selection = None
        self.augmentation = augmentation

    def mask_one_cdr_(self, data, cdr_to_mask):
        cdr_flag = data['cdr_flag']

        cdr_to_mask_flag = (cdr_flag == cdr_to_mask)
        if self.augmentation:
            cdr_to_mask_flag = random_shrink_extend(cdr_to_mask_flag)

        cdr_first, cdr_last = continuous_flag_to_range(cdr_to_mask_flag)
        left_idx = max(0, cdr_first-1)
        right_idx = min(data['aa'].size(0)-1, cdr_last+1)
        anchor_flag = torch.zeros(data['aa'].shape, dtype=torch.bool)
        anchor_flag[left_idx] = True
        anchor_flag[right_idx] = True

        if 'generate_flag' not in data:
            data['generate_flag'] = cdr_to_mask_flag
            data['anchor_flag'] = anchor_flag
        else:
            data['generate_flag'] |= cdr_to_mask_flag
            data['anchor_flag'] |= anchor_flag

    def mask_for_one_chain_(self, data):
        cdr_flag = data['cdr_flag']
        cdr_all = cdr_flag[cdr_flag > 0].unique().tolist()
    
        num_cdrs_to_mask = random.randint(1, len(cdr_all))

        if self.selection is not None:
            cdrs_to_mask = list(set(cdr_all).intersection(self.selection))
        else:
            random.shuffle(cdr_all)
            cdrs_to_mask = cdr_all[:num_cdrs_to_mask]

        for cdr_to_mask in cdrs_to_mask:
            self.mask_one_cdr_(data, cdr_to_mask)

    def __call__(self, structure):
        if structure['heavy'] is not None:
            self.mask_for_one_chain_(structure['heavy'])
        if structure['light'] is not None:
            self.mask_for_one_chain_(structure['light'])
        return structure

        
@register_transform('mask_antibody')
class MaskAntibody(object):

    def mask_ab_chain_(self, data):
        data['generate_flag'] = torch.ones(data['aa'].shape, dtype=torch.bool)

    def __call__(self, structure):
        pos_ab_alpha = []
        if structure['heavy'] is not None:
            self.mask_ab_chain_(structure['heavy'])
            pos_ab_alpha.append(
                structure['heavy']['pos_heavyatom'][:, constants.BBHeavyAtom.CA]
            )
        if structure['light'] is not None:
            self.mask_ab_chain_(structure['light'])
            pos_ab_alpha.append(
                structure['light']['pos_heavyatom'][:, constants.BBHeavyAtom.CA]
            )
        pos_ab_alpha = torch.cat(pos_ab_alpha, dim=0)   # (L_Ab, 3)

        if structure['antigen'] is not None:
            pos_ag_alpha = structure['antigen']['pos_heavyatom'][:, constants.BBHeavyAtom.CA]
            ag_ab_dist = torch.cdist(pos_ag_alpha, pos_ab_alpha)    # (L_Ag, L_Ab)
            nn_ab_dist = ag_ab_dist.min(dim=1)[0]   # (L_Ag)
            contact_flag = (nn_ab_dist <= 6.0)      # (L_Ag)
            if contact_flag.sum().item() == 0:
                contact_flag[nn_ab_dist.argmin()] = True

            anchor_idx = torch.multinomial(contact_flag.float(), num_samples=1).item()
            anchor_flag = torch.zeros(structure['antigen']['aa'].shape, dtype=torch.bool)
            anchor_flag[anchor_idx] = True
            structure['antigen']['anchor_flag'] = anchor_flag
            structure['antigen']['contact_flag'] = contact_flag
        
        return structure


@register_transform('remove_antigen')
class RemoveAntigen:

    def __call__(self, structure):
        structure['antigen'] = None
        structure['antigen_seqmap'] = None
        return structure

AbDesign/diffab/utils/transforms/mask.py
pengzhangzhi/ab_opt
import torch

from ..protein import constants
from ._base import register_transform


@register_transform('merge_chains')
class MergeChains(object):

    def __init__(self):
        super().__init__()

    def assign_chain_number_(self, data_list):
        chains = set()
        for data in data_list:
            chains.update(data['chain_id'])
        chains = {c: i for i, c in enumerate(chains)}

        for data in data_list:
            data['chain_nb'] = torch.LongTensor([
                chains[c] for c in data['chain_id']
            ])

    def _data_attr(self, data, name):
        if name in ('generate_flag', 'anchor_flag') and name not in data:
            return torch.zeros(data['aa'].shape, dtype=torch.bool)
        else:
            return data[name]

    def __call__(self, structure):
        data_list = []
        if structure['heavy'] is not None:
            structure['heavy']['fragment_type'] = torch.full_like(
                structure['heavy']['aa'],
                fill_value = constants.Fragment.Heavy,
            )
            data_list.append(structure['heavy'])

        if structure['light'] is not None:
            structure['light']['fragment_type'] = torch.full_like(
                structure['light']['aa'],
                fill_value = constants.Fragment.Light,
            )
            data_list.append(structure['light'])

        if structure['antigen'] is not None:
            structure['antigen']['fragment_type'] = torch.full_like(
                structure['antigen']['aa'],
                fill_value = constants.Fragment.Antigen,
            )
            structure['antigen']['cdr_flag'] = torch.zeros_like(
                structure['antigen']['aa'],
            )
            data_list.append(structure['antigen'])

        self.assign_chain_number_(data_list)

        list_props = {
            'chain_id': [],
            'icode': [],
        }
        tensor_props = {
            'chain_nb': [],
            'resseq': [],
            'res_nb': [],
            'aa': [],
            'pos_heavyatom': [],
            'mask_heavyatom': [],
            'generate_flag': [],
            'cdr_flag': [],
            'anchor_flag': [],
            'fragment_type': [],
        }

        for data in data_list:
            for k in list_props.keys():
                list_props[k].append(self._data_attr(data, k))
            for k in tensor_props.keys():
                tensor_props[k].append(self._data_attr(data, k))

        list_props = {k: sum(v, start=[]) for k, v in list_props.items()}
        tensor_props = {k: torch.cat(v, dim=0) for k, v in tensor_props.items()}
        data_out = {
            **list_props,
            **tensor_props,
        }
        return data_out


AbDesign/diffab/utils/transforms/merge.py
pengzhangzhi/ab_opt
import torch

from ._base import _mask_select_data, register_transform
from ..protein import constants

class CUTOFFVIOLATION(Exception):
    pass

@register_transform('patch_around_anchor')
class PatchAroundAnchor(object):

    def __init__(self, initial_patch_size=128, antigen_size=128,max_nb_hotspots=3,is_train=True, dist_cutoff=None,max_distance=40):
        super().__init__()
        self.initial_patch_size = initial_patch_size
        self.antigen_size = antigen_size
        self.max_nb_hotspots = max_nb_hotspots
        self.is_train = is_train
        self.dist_cutoff = dist_cutoff
        self.max_distance = max_distance
        
    def _center(self, data, origin):
        origin = origin.reshape(1, 1, 3)
        data['pos_heavyatom'] -= origin # (L, A, 3)
        data['pos_heavyatom'] = data['pos_heavyatom'] * data['mask_heavyatom'][:, :, None]
        data['origin'] = origin.reshape(3)
        return data

    def __call__(self, data):
        anchor_flag = data['anchor_flag']   # (L,)
        anchor_points = data['pos_heavyatom'][anchor_flag, constants.BBHeavyAtom.CA]    # (n_anchors, 3)
        antigen_mask = (data['fragment_type'] == constants.Fragment.Antigen)
        antibody_mask = torch.logical_not(antigen_mask)

        if anchor_flag.sum().item() == 0:
            # Generating full antibody-Fv, no antigen given
            data_patch = _mask_select_data(
                data = data,
                mask = antibody_mask,
            )
            data_patch = self._center(
                data_patch,
                origin = data_patch['pos_heavyatom'][:, constants.BBHeavyAtom.CA].mean(dim=0)
            )
            return data_patch

        pos_alpha = data['pos_heavyatom'][:, constants.BBHeavyAtom.CA]  # (L, 3)
        dist_anchor = torch.cdist(pos_alpha, anchor_points).min(dim=1)[0]    # (L, )
        initial_patch_idx = torch.topk(
            dist_anchor,
            k = min(self.initial_patch_size, dist_anchor.size(0)),
            largest=False,
        )[1]   # (initial_patch_size, )

        dist_anchor_antigen = dist_anchor.masked_fill(
            mask = antibody_mask, # Fill antibody with +inf
            value = float('+inf')
        )   # (L, )
        antigen_dist, antigen_patch_idx = torch.topk(
            dist_anchor_antigen, 
            k = min(self.antigen_size, antigen_mask.sum().item()), 
            largest=False, sorted=True
        )    # (ag_size, )
        
        # label the nb_hotspots closest antigen residues as the hotspots 
        # nb_hotspots = int(torch.randint(0, self.max_nb_hotspots+1, (1, ))) if self.is_train else self.max_nb_hotspots
        nb_hotspots = 1 if self.max_nb_hotspots == 0 else self.max_nb_hotspots
        nb_hotspots = min(nb_hotspots, antigen_mask.sum().item())
        dist_cdr = torch.cdist(
            pos_alpha,
            pos_alpha[data['generate_flag']]
        ).min(dim=1)[0]
        dist_cdr = dist_cdr.masked_fill(
            mask = antibody_mask,
            value = float('+inf')
        )
        dist_cdr, cdr2antigen_min_idx = torch.topk(
            dist_cdr, 
            k = nb_hotspots,
            largest=False, sorted=True
        )    # (nb_hotspots, )
        if self.dist_cutoff is not None and dist_cdr.min() > self.dist_cutoff:
            raise CUTOFFVIOLATION(f'No interaction pairs found within cutoff of {self.dist_cutoff}')
        # antibody-antigen interaction info 
        # the antigen residues that has cloest distance to the antibody residues
        data['antigen_hotspots'] = torch.ones_like(data['generate_flag'], dtype=torch.int64) 
        data['antigen_hotspots'][cdr2antigen_min_idx] = 2 
        
        # the cloest distance
        data['to_hotspot_dist'] = torch.zeros_like(data['generate_flag'], dtype=torch.float32)
        data['to_hotspot_dist'][cdr2antigen_min_idx] = dist_cdr
        
        hotspot_label = torch.ones_like(data['generate_flag'], dtype=torch.int64)
        hotspot_distance = torch.ones_like(data['generate_flag'], dtype=torch.long)
        if self.max_nb_hotspots > 0:
            # 2 for hotspot, 1 for non-hotspot. 0 saved for padding
            hotspot_label[cdr2antigen_min_idx] = 2 
            hotspot_distance[cdr2antigen_min_idx] = torch.floor(dist_cdr).long().clamp(1, self.max_distance-1)
        data['hotspot_label'] = hotspot_label
        data['hotspot_distance'] = hotspot_distance
        ### patch ###
        patch_mask = torch.logical_or(
            data['generate_flag'],
            data['anchor_flag'],
        )
        patch_mask[initial_patch_idx] = True
        patch_mask[antigen_patch_idx] = True

        patch_idx = torch.arange(0, patch_mask.shape[0])[patch_mask]

        data_patch = _mask_select_data(data, patch_mask)
        data_patch = self._center(
            data_patch,
            origin = anchor_points.mean(dim=0)
        )
        data_patch['patch_idx'] = patch_idx
        
        return data_patch

AbDesign/diffab/utils/transforms/patch.py
pengzhangzhi/ab_opt

from ._base import register_transform


@register_transform('select_atom')
class SelectAtom(object):

    def __init__(self, resolution):
        super().__init__()
        assert resolution in ('full', 'backbone')
        self.resolution = resolution

    def __call__(self, data):
        if self.resolution == 'full':
            data['pos_atoms'] = data['pos_heavyatom'][:, :]
            data['mask_atoms'] = data['mask_heavyatom'][:, :]
        elif self.resolution == 'backbone':
            data['pos_atoms'] = data['pos_heavyatom'][:, :5]
            data['mask_atoms'] = data['mask_heavyatom'][:, :5]
        return data

AbDesign/diffab/utils/transforms/select_atom.py
pengzhangzhi/ab_opt
from collections import defaultdict
import os
import argparse
import copy
import json
from tqdm.auto import tqdm
from torch.utils.data import DataLoader

from diffab.datasets import get_dataset
from diffab.models import get_model
from diffab.modules.common.geometry import reconstruct_backbone_partially
from diffab.modules.common.so3 import so3vec_to_rotation
from diffab.utils.inference import RemoveNative
from diffab.utils.protein.constants import BBHeavyAtom
from diffab.utils.protein.writers import save_pdb
from diffab.utils.misc import *
from diffab.utils.data import *
from diffab.utils.transforms import *
from diffab.utils.inference import *


def recursive_to(obj, device):
    if isinstance(obj, torch.Tensor):
        if device == 'cpu':
            return obj.cpu()
        try:
            return obj.cuda(device=device, non_blocking=True)
        except RuntimeError:
            return obj.to(device)
    elif isinstance(obj, list):
        return [recursive_to(o, device=device) for o in obj]
    elif isinstance(obj, tuple):
        return tuple(recursive_to(o, device=device) for o in obj)
    elif isinstance(obj, dict):
        return {k: recursive_to(v, device=device) for k, v in obj.items()}

    else:
        return obj


def create_data_variants(config, structure_factory):
    structure = structure_factory()
    structure_id = structure['id']

    data_variants = []
    if config.mode == 'single_cdr':
        cdrs = sorted(list(set(find_cdrs(structure)).intersection(config.sampling.cdrs)))
        for cdr_name in cdrs:
            transform = Compose([
                MaskSingleCDR(cdr_name, augmentation=False),
                MergeChains(),
            ])
            data_var = transform(structure_factory())
            residue_first, residue_last = get_residue_first_last(data_var)
            data_variants.append({
                'data': data_var,
                'name': f'{structure_id}-{cdr_name}',
                'tag': f'{cdr_name}',
                'cdr': cdr_name,
                'residue_first': residue_first,
                'residue_last': residue_last,
            })
    elif config.mode == 'multiple_cdrs':
        cdrs = sorted(list(set(find_cdrs(structure)).intersection(config.sampling.cdrs)))
        transform = Compose([
            MaskMultipleCDRs(selection=cdrs, augmentation=False),
            MergeChains(),
        ])
        data_var = transform(structure_factory())
        data_variants.append({
            'data': data_var,
            'name': f'{structure_id}-MultipleCDRs',
            'tag': 'MultipleCDRs',
            'cdrs': cdrs,
            'residue_first': None,
            'residue_last': None,
        })
    elif config.mode == 'full':
        transform = Compose([
            MaskAntibody(),
            MergeChains(),
        ])
        data_var = transform(structure_factory())
        data_variants.append({
            'data': data_var,
            'name': f'{structure_id}-Full',
            'tag': 'Full',
            'residue_first': None,
            'residue_last': None,
        })
    elif config.mode == 'abopt':
        cdrs = sorted(list(set(find_cdrs(structure)).intersection(config.sampling.cdrs)))
        for cdr_name in cdrs:
            transform = Compose([
                MaskSingleCDR(cdr_name, augmentation=False),
                MergeChains(),
            ])
            data_var = transform(structure_factory())
            residue_first, residue_last = get_residue_first_last(data_var)
            for opt_step in config.sampling.optimize_steps:
                data_variants.append({
                    'data': data_var,
                    'name': f'{structure_id}-{cdr_name}-O{opt_step}',
                    'tag': f'{cdr_name}-O{opt_step}',
                    'cdr': cdr_name,
                    'opt_step': opt_step,
                    'residue_first': residue_first,
                    'residue_last': residue_last,
                })
    else:
        raise ValueError(f'Unknown mode: {config.mode}.')
    return data_variants


def run_on_variant(config, log_dir, logger, data_native, model, inference_tfm, variant,**kwargs):
    """
    return:
        result_dict: dict, keys are rmsd, aa_recovery, aa_seq, value is lists. 
    """
    
    os.makedirs(os.path.join(log_dir, variant['tag']), exist_ok=True)
    logger.info(f"Start sampling for: {variant['tag']}")

    save_pdb(data_native, os.path.join(log_dir, variant['tag'], 'REF1.pdb'))       # w/  OpenMM minimization
    
    data_cropped = inference_tfm(
            copy.deepcopy(variant['data'])
    )
    batch_size = min(100, config.sampling.num_samples)
    data_list_repeat = [ data_cropped ] * config.sampling.num_samples
    collate_fn = CollateNodes(eight=False)
    loader = DataLoader(data_list_repeat, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
        
    count = 0
    result_dict = defaultdict(list)
    for batch in tqdm(loader, desc=variant['name'], dynamic_ncols=True):
        # batch size
        count, result = run_on_batch(config, log_dir, model, variant, data_cropped, count, batch,**kwargs)
        count += batch_size
        for k, v in result.items():
            result_dict[k].extend(v)
        
    logger.info('Finished.\n')
    return result_dict


def run_on_batch(config, log_dir, model, variant, data_cropped, count, batch, traj_idx = 0):
    torch.set_grad_enabled(False)
    model.eval()
    device = next(model.parameters()).device
    batch = recursive_to(batch, device)
    if 'abopt' in config.mode:
                # Antibody optimization starting from native
        traj_batch = model.optimize(batch, opt_step=variant['opt_step'], optimize_opt={
                    'pbar': True,
                    'sample_structure': config.sampling.sample_structure,
                    'sample_sequence': config.sampling.sample_sequence,
                })
    else:
                # De novo design
        traj_batch = model.sample(batch,
                sample_opt={
                    'pbar': True,
                    'sample_structure': config.sampling.sample_structure,
                    'sample_sequence': config.sampling.sample_sequence,
                },
                
                )
    
    result_dict = save_trajectory(log_dir, variant, data_cropped, count, batch, traj_idx, device, traj_batch)
    return count, result_dict

def save_trajectory(log_dir, variant, data_cropped, count, batch, traj_idx, device, traj_batch):
    """
    return:
        result_dict: dict, keys are rmsd, aa_recovery, aa_seq, value is lists. 
    """
    r_new = so3vec_to_rotation(traj_batch[traj_idx][0]).to(device)
    t_new = (traj_batch[traj_idx][1]).to(device)
    aa_new = (traj_batch[traj_idx][2]).to(device)   # 2: Amino acid.
    pos_atom_new, mask_atom_new = reconstruct_backbone_partially(
                pos_ctx = batch['pos_heavyatom'],
                R_new = r_new,
                t_new = t_new,
                aa = aa_new,
                chain_nb = batch['chain_nb'],
                res_nb = batch['res_nb'],
                mask_atoms = batch['mask_heavyatom'],
                mask_recons = batch['generate_flag'],
            )
    aa_new = aa_new.cpu()
    pos_atom_new = pos_atom_new.cpu()
    mask_atom_new = mask_atom_new.cpu()
    result_dict = defaultdict(list)
    for i in range(aa_new.size(0)):
        result,pdb_path = save_each_sample(log_dir, variant, data_cropped, count, batch, aa_new, pos_atom_new, mask_atom_new, i)
        count += 1
        # append items in result to result_dict
        for k, v in result.items():
            result_dict[k].append(v)
    return result_dict



AbDesign/diffab/utils/val.py
pengzhangzhi/ab_opt
from diffab.tools.relax.run import main

if __name__ == '__main__':
    main()

AbDesign/relax.py
pengzhangzhi/ab_opt
import os
import shutil
import argparse
import torch
import glob
# import torch.utils.tensorboard
from torch.nn.utils import clip_grad_norm_
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

from diffab.datasets import get_dataset
from diffab.models import get_model
from diffab.utils.misc import *
from diffab.utils.data import *
from diffab.utils.train import *


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='configs/train/codesign_single.yml')
    parser.add_argument('--logdir', type=str, default='./logs')
    parser.add_argument('--debug', action='store_true', default=False)
    parser.add_argument('--device', type=str, default='cuda')
    parser.add_argument('--num_workers', type=int, default=8)
    parser.add_argument('--tag', type=str, default='')
    parser.add_argument('--resume', type=str, default=None, help='ckpt path. Resume from a checkpoint.')
    parser.add_argument('--finetune', type=str, default=None, help='ckpt path. Finetune from a pretrained model with extra finetune_iters.')
    parser.add_argument('--finetune_iters', type=int, default=200_000)
    args = parser.parse_args()

    # Load configs
    config, config_name = load_config(args.config)
    seed_all(config.train.seed)

    if os.environ.get('wandb') is not None:
        import wandb
        wandb.init(project="diffusion-antibody", config=dict(config), name=f"{args.tag}-{config_name}")
    
    
    # Logging
    if args.debug:
        logger = get_logger('train', None)
        writer = BlackHole()
    else:
        if args.resume:
            log_dir = os.path.dirname(os.path.dirname(args.resume))
        else:
            log_dir = get_new_log_dir(args.logdir, prefix=config_name, tag=args.tag)
        ckpt_dir = os.path.join(log_dir, 'checkpoints')
        if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir)
        logger = get_logger('train', log_dir)
        writer = BlackHole() # torch.utils.tensorboard.SummaryWriter(log_dir)
        # tensorboard_trace_handler = torch.profiler.tensorboard_trace_handler(log_dir)
        if not os.path.exists(os.path.join(log_dir, os.path.basename(args.config))):
            shutil.copyfile(args.config, os.path.join(log_dir, os.path.basename(args.config)))
    logger.info(args)
    logger.info(config)

    # Data
    logger.info('Loading dataset...')
    train_dataset = get_dataset(config.dataset.train)
    val_dataset = get_dataset(config.dataset.val)
    train_iterator = inf_iterator(DataLoader(
        train_dataset, 
        batch_size=config.train.batch_size, 
        collate_fn=CollateNodes(),
        shuffle=True,
        num_workers=args.num_workers
    ))
    val_loader = DataLoader(val_dataset, batch_size=config.train.batch_size, collate_fn=CollateNodes(), shuffle=False, num_workers=args.num_workers)
    logger.info('Train %d | Val %d' % (len(train_dataset), len(val_dataset)))

    # Model
    logger.info('Building model...')
    model = get_model(config.model).to(args.device)
    logger.info('Number of parameters: %d' % count_parameters(model))

    # Optimizer & scheduler
    optimizer = get_optimizer(config.train.optimizer, model)
    scheduler = get_scheduler(config.train.scheduler, optimizer)
    optimizer.zero_grad()
    it_first = 1

    # Resume
    if args.resume is not None or args.finetune is not None:
        ckpt_path = args.resume if args.resume is not None else args.finetune
        logger.info('Resuming from checkpoint: %s' % ckpt_path)
        ckpt = torch.load(ckpt_path, map_location=args.device)
        it_first = ckpt['iteration']  # + 1
        model.load_state_dict(ckpt['model'])
        logger.info('Resuming optimizer states...')
        optimizer.load_state_dict(ckpt['optimizer'])
        logger.info('Resuming scheduler states...')
        scheduler.load_state_dict(ckpt['scheduler'])
        

    # Train
    def train(it):
        time_start = current_milli_time()
        model.train()

        # Prepare data
        batch = recursive_to(next(train_iterator), args.device)

        # Forward
        # if args.debug: torch.set_anomaly_enabled(True)
        loss_dict = model(batch)
        loss = sum_weighted_losses(loss_dict, config.train.loss_weights)
        loss_dict['overall'] = loss
        time_forward_end = current_milli_time()

        # Backward
        loss.backward()
        orig_grad_norm = clip_grad_norm_(model.parameters(), config.train.max_grad_norm)
        optimizer.step()
        optimizer.zero_grad()
        time_backward_end = current_milli_time()

        # Logging
        log_losses(loss_dict, it, 'train', logger, writer, others={
            'grad': orig_grad_norm,
            'lr': optimizer.param_groups[0]['lr'],
            'time_forward': (time_forward_end - time_start) / 1000,
            'time_backward': (time_backward_end - time_forward_end) / 1000,
        })
        if os.environ.get('wandb') is not None:
            track_dict = {
                **{k: float(v) for k, v in loss_dict.items()},
                **{
                    "grad": orig_grad_norm.item(),
                    "lr": optimizer.param_groups[0]["lr"],
                },
            }
            wandb.log(track_dict)
            
        if not torch.isfinite(loss):
            logger.error('NaN or Inf detected.')
            torch.save({
                'config': config,
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'scheduler': scheduler.state_dict(),
                'iteration': it,
                'batch': recursive_to(batch, 'cpu'),
            }, os.path.join(log_dir, 'checkpoint_nan_%d.pt' % it))
            raise KeyboardInterrupt()

    # Validate
    def validate(it):
        save_dir = os.path.join(log_dir, 'val', str(it))
        os.makedirs(save_dir, exist_ok=True)
        val_result = eval_on_dataset(config, val_dataset, model, logger, save_dir)
        # find a file ends with pdb in the subdirectory of save_dir recursively
        if os.environ.get('wandb') is not None:
            wandb.log(val_result)
            pdb_file = glob(os.path.join(save_dir, '*', '*', '*patch.pdb'))[0]
            wandb.log({'val_sample': wandb.Molecule(pdb_file)})
            
        return val_result
    max_iters = it_first+args.finetune_iters if args.finetune is not None else config.train.max_iters
    min_rmsd = 1e6
    try:
        # validate(0)
        for it in range(it_first, max_iters + 1):
            train(it)
            if it % config.train.val_freq == 0:
                val_result = validate(it)
                if not args.debug:
                    if val_result['H_CDR3-rmsd'] < min_rmsd:
                        min_rmsd = val_result['H_CDR3-rmsd']
                        aar = val_result['H_CDR3-aa_recovery']
                        fname = f'best-{it}-RMSD-{float(min_rmsd):.3f}-AAR-{float(aar):.3f}.pt'
                    else: fname = '%d.pt' % it
                    ckpt_path = os.path.join(ckpt_dir, fname)
                    torch.save({
                        'config': config,
                        'model': model.state_dict(),
                        'optimizer': optimizer.state_dict(),
                        'scheduler': scheduler.state_dict(),
                        'iteration': it,
                        'val_result': val_result,
                    }, ckpt_path)
    except KeyboardInterrupt:
        logger.info('Terminating...')

AbDesign/train.py
pengzhangzhi/ab_opt
#!/usr/bin/env python

import Bio.PDB
import warnings
from Bio import BiopythonWarning
warnings.simplefilter('ignore', BiopythonWarning)
import sys
import os
import re
import tempfile
import numpy as np
from Bio.SVDSuperimposer import SVDSuperimposer
from math import sqrt
from argparse import ArgumentParser
import itertools
import subprocess

def parse_fnat(fnat_out):
    fnat=-1;
    nat_correct=-1
    nat_total=-1
    fnonnat=-1
    nonnat_count=-1
    model_total=-1
    inter=[]
    for line in fnat_out.split("\n"):
#        print line
        line=line.rstrip('\n')
        match=re.search(r'NATIVE: (\d+)(\w) (\d+)(\w)',line)
        if(re.search(r'^Fnat',line)):
            list=line.split(' ')
            fnat=float(list[3])
            nat_correct=int(list[1])
            nat_total=int(list[2])
        elif(re.search(r'^Fnonnat',line)):
            list=line.split(' ')
            fnonnat=float(list[3])
            nonnat_count=int(list[1])
            model_total=int(list[2])
        elif(match):
            #print line
            res1=match.group(1)
            chain1=match.group(2)
            res2=match.group(3)
            chain2=match.group(4)
           # print res1 + ' ' + chain1 + ' ' + res2 + ' ' + chain2
            inter.append(res1 + chain1)
            inter.append(res2 + chain2)
    return (fnat,nat_correct,nat_total,fnonnat,nonnat_count,model_total,inter)

def capri_class(fnat,iRMS,LRMS,capri_peptide=False):


    if capri_peptide:
               
        if(fnat < 0.2 or (LRMS > 5.0 and iRMS > 2.0)):
            return 'Incorrect'
        elif((fnat >= 0.2 and fnat < 0.5) and (LRMS <= 5.0 or iRMS <= 2.0) or (fnat >= 0.5 and LRMS > 2.0 and iRMS > 1.0)):
            return 'Acceptable'
        elif((fnat >= 0.5 and fnat < 0.8) and (LRMS <= 2.0 or iRMS <= 1.0) or (fnat >= 0.8 and LRMS > 1.0 and iRMS > 0.5)):
            return 'Medium'
        elif(fnat >= 0.8 and (LRMS <= 1.0 or iRMS <= 0.5)):
            return 'High'
        else:
            return 'Undef'
    else:

        if(fnat < 0.1 or (LRMS > 10.0 and iRMS > 4.0)):
            return 'Incorrect'
        elif((fnat >= 0.1 and fnat < 0.3) and (LRMS <= 10.0 or iRMS <= 4.0) or (fnat >= 0.3 and LRMS > 5.0 and iRMS > 2.0)):
            return 'Acceptable'
        elif((fnat >= 0.3 and fnat < 0.5) and (LRMS <= 5.0 or iRMS <= 2.0) or (fnat >= 0.5 and LRMS > 1.0 and iRMS > 1.0)):
            return 'Medium'
        elif(fnat >= 0.5 and (LRMS <= 1.0 or iRMS <= 1.0)):
            return 'High'
        else:
            return 'Undef'


def capri_class_DockQ(DockQ,capri_peptide=False):

    if capri_peptide:
        return 'Undef for capri_peptides'
    
    (c1,c2,c3)=(0.23,0.49,0.80)
    if(DockQ < c1):
        return 'Incorrect'
    elif(DockQ >= c1 and DockQ < c2):
        return 'Acceptable'
    elif(DockQ >= c2 and DockQ < c3):
        return 'Medium'
    elif(DockQ >= c3):
        return 'High'
    else:
        return 'Undef'


def calc_DockQ(model,native,use_CA_only=False,capri_peptide=False):
    
#    exec_path=os.path.dirname(os.path.abspath(sys.argv[0]))    
    exec_path=os.path.dirname(os.path.abspath(__file__))
    atom_for_sup=['CA','C','N','O']
    if(use_CA_only):
        atom_for_sup=['CA']


    cmd_fnat=exec_path + '/fnat ' + model + ' ' + native + ' 5 -all'
    cmd_interface=exec_path + '/fnat ' + model + ' ' + native + ' 10 -all'

    if capri_peptide:
        cmd_fnat=exec_path + '/fnat ' + model + ' ' + native + ' 4 -all'
        cmd_interface=exec_path + '/fnat ' + model + ' ' + native + ' 8 -cb' 

    fnat_out = os.popen(cmd_fnat).read()

    #fnat_out = subprocess.getoutput(cmd_fnat)
    #print(fnat_out)
    #    sys.exit()
    (fnat,nat_correct,nat_total,fnonnat,nonnat_count,model_total,interface5A)=parse_fnat(fnat_out)
    assert fnat!=-1, "Error running cmd: %s\n" % (cmd_fnat)
    inter_out = os.popen(cmd_interface).read()
#   inter_out = subprocess.getoutput(cmd_interface)
    
    (fnat_bb,nat_correct_bb,nat_total_bb,fnonnat_bb,nonnat_count_bb,model_total_bb,interface)=parse_fnat(inter_out)
    assert fnat_bb!=-1, "Error running cmd: %s\n" % (cmd_interface)

    #print fnat
    #Use same interface as for fnat for iRMS
    #interface=interface5A


    # Start the parser
    pdb_parser = Bio.PDB.PDBParser(QUIET = True)

    # Get the structures
    ref_structure = pdb_parser.get_structure("reference", native)
    sample_structure = pdb_parser.get_structure("model", model)

    # Use the first model in the pdb-files for alignment
    # Change the number 0 if you want to align to another structure
    ref_model    = ref_structure[0]
    sample_model = sample_structure[0]

    # Make a list of the atoms (in the structures) you wish to align.
    # In this case we use CA atoms whose index is in the specified range
    ref_atoms = []
    sample_atoms = []

    common_interface=[]

    chain_res={}


    #find atoms common in both sample and native
    atoms_def_sample=[]
    atoms_def_in_both=[]
    #first read in sample
    for sample_chain in sample_model:
#        print sample_chain
        chain=sample_chain.id
#        print chain
        for sample_res in sample_chain:
           # print sample_res
            if sample_res.get_id()[0] != ' ': #Skip hetatm.
                continue
            resname=sample_res.get_id()[1]
            key=str(resname) + chain
            for a in atom_for_sup:
                atom_key=key + '.' + a
                if a in sample_res:
                    if atom_key in atoms_def_sample:
                        ...
                        # print(atom_key + ' already added (MODEL)!!!')
                    atoms_def_sample.append(atom_key)

    #then read in native also present in sample
    for ref_chain in ref_model:
        chain=ref_chain.id
        for ref_res in ref_chain:
            #print ref_res
            if ref_res.get_id()[0] != ' ': #Skip hetatm.
#                print ref_res.get_id()
                continue
            resname=ref_res.get_id()[1]
            key=str(resname) + chain
            for a in atom_for_sup:
                atom_key=key + '.' + a
                if a in ref_res and atom_key in atoms_def_sample:
                    if atom_key in atoms_def_in_both:
                        ...
                        # print(atom_key + ' already added (Native)!!!') 
                    atoms_def_in_both.append(atom_key)


#    print atoms_def_in_both
    for sample_chain in sample_model:
        chain=sample_chain.id
        if chain not in list(chain_res.keys()):
            chain_res[chain]=[]
        for sample_res in sample_chain:
            if sample_res.get_id()[0] != ' ': #Skip hetatm.
                continue
            resname=sample_res.get_id()[1]
            key=str(resname) + chain
            chain_res[chain].append(key)
            if key in interface:
                for a in atom_for_sup:
                    atom_key=key + '.' + a
                    if a in sample_res and atom_key in atoms_def_in_both:
                        sample_atoms.append(sample_res[a])
                common_interface.append(key)

    #print inter_pairs

    chain_ref={}
    common_residues=[]



    # Iterate of all chains in the model in order to find all residues
    for ref_chain in ref_model:
        # Iterate of all residues in each model in order to find proper atoms
        #  print dir(ref_chain)
        chain=ref_chain.id
        if chain not in list(chain_ref.keys()):
            chain_ref[chain]=[]
        for ref_res in ref_chain:
            if ref_res.get_id()[0] != ' ': #Skip hetatm.
                continue
            resname=ref_res.get_id()[1]
            key=str(resname) + chain

            #print ref_res
            #      print key
            # print chain_res.values()
            if key in chain_res[chain]: # if key is present in sample
                #print key
                for a in atom_for_sup:
                    atom_key=key + '.' + a
                    if a in ref_res and atom_key in atoms_def_in_both:
                        chain_ref[chain].append(ref_res[a])
                        common_residues.append(key)
                      #chain_sample.append((ref_res['CA'])
            if key in common_interface:
              # Check if residue number ( .get_id() ) is in the list
              # Append CA atom to list
                #print key  
                for a in atom_for_sup:
                    atom_key=key + '.' + a
                    #print atom_key
                    if a in ref_res and atom_key in atoms_def_in_both:
                        ref_atoms.append(ref_res[a])



    #get the ones that are present in native        
    chain_sample={}
    for sample_chain in sample_model:
        chain=sample_chain.id
        if chain not in list(chain_sample.keys()):
            chain_sample[chain]=[]
        for sample_res in sample_chain:
            if sample_res.get_id()[0] != ' ': #Skip hetatm.
                continue
            resname=sample_res.get_id()[1]
            key=str(resname) + chain
            if key in common_residues:
                for a in atom_for_sup:
                    atom_key=key + '.' + a
                    if a in sample_res and atom_key in atoms_def_in_both:
                        chain_sample[chain].append(sample_res[a])

        #if key in common_residues:
        #     print key  
        #sample_atoms.append(sample_res['CA'])
        #common_interface.append(key)

        
    assert len(ref_atoms)!=0, "length of native is zero"
    assert len(sample_atoms)!=0, "length of model is zero"
    assert len(ref_atoms)==len(sample_atoms), "Different number of atoms in native and model %d %d\n" % (len(ref_atoms),len(sample_atoms))

    super_imposer = Bio.PDB.Superimposer()
    super_imposer.set_atoms(ref_atoms, sample_atoms)
    super_imposer.apply(sample_model.get_atoms())

    # Print RMSD:
    irms=super_imposer.rms

    (chain1,chain2)=list(chain_sample.keys())

    ligand_chain=chain1
    receptor_chain=chain2
    len1=len(chain_res[chain1])
    len2=len(chain_res[chain2])

    assert len1!=0, "%s chain has zero length!\n" % chain1
    assert len2!=0, "%s chain has zero length!\n" % chain2

    class1='ligand'
    class2='receptor'
    if(len(chain_sample[chain1]) > len(chain_sample[chain2])):
        receptor_chain=chain1
        ligand_chain=chain2
        class1='receptor'
        class2='ligand'



    #print len1
    #print len2
    #print chain_sample.keys()

    #Set to align on receptor
    assert len(chain_ref[receptor_chain])==len(chain_sample[receptor_chain]), "Different number of atoms in native and model receptor (chain %c) %d %d\n" % (receptor_chain,len(chain_ref[receptor_chain]),len(chain_sample[receptor_chain]))

    super_imposer.set_atoms(chain_ref[receptor_chain], chain_sample[receptor_chain])
    super_imposer.apply(sample_model.get_atoms())
    receptor_chain_rms=super_imposer.rms
    #print receptor_chain_rms
    #print dir(super_imposer)
    #print chain1_rms

    #Grep out the transformed ligand coords

    #print ligand_chain

    #print chain_ref[ligand_chain]
    #print chain_sample[ligand_chain]
    #l1=len(chain_ref[ligand_chain])
    #l2=len(chain_sample[ligand_chain])




    assert len(chain_ref[ligand_chain])!=0 or len(chain_sample[ligand_chain])!=0, "Zero number of equivalent atoms in native and model ligand (chain %s) %d %d.\nCheck that the residue numbers in model and native is consistent\n" % (ligand_chain,len(chain_ref[ligand_chain]),len(chain_sample[ligand_chain]))


    assert len(chain_ref[ligand_chain])==len(chain_sample[ligand_chain]), "Different number of atoms in native and model ligand (chain %c) %d %d\n" % (ligand_chain,len(chain_ref[ligand_chain]),len(chain_sample[ligand_chain]))

    coord1=np.array([atom.coord for atom in chain_ref[ligand_chain]])
    coord2=np.array([atom.coord for atom in chain_sample[ligand_chain]])

    #coord1=np.array([atom.coord for atom in chain_ref[receptor_chain]])
    #coord2=np.array([atom.coord for atom in chain_sample[receptor_chain]])

    #print len(coord1)
    #print len(coord2)

    sup=SVDSuperimposer()
    Lrms = sup._rms(coord1,coord2) #using the private _rms function which does not superimpose


    #super_imposer.set_atoms(chain_ref[ligand_chain], chain_sample[ligand_chain])
    #super_imposer.apply(sample_model.get_atoms())
    #coord1=np.array([atom.coord for atom in chain_ref[receptor_chain]])
    #coord2=np.array([atom.coord for atom in chain_sample[receptor_chain]])
    #Rrms= sup._rms(coord1,coord2)
    #should give same result as above line
    #diff = coord1-coord2
    #l = len(diff) #number of atoms
    #from math import sqrt
    #print sqrt(sum(sum(diff*diff))/l)
    #print np.sqrt(np.sum(diff**2)/l)
    DockQ=(float(fnat) + 1/(1+(irms/1.5)*(irms/1.5)) + 1/(1+(Lrms/8.5)*(Lrms/8.5)))/3
    info={}
    info['DockQ']=DockQ
    info['irms']=irms
    info['Lrms']=Lrms
    info['fnat']=fnat
    info['nat_correct']=nat_correct
    info['nat_total']=nat_total

    info['fnonnat']=fnonnat
    info['nonnat_count']=nonnat_count
    info['model_total']=model_total
    
    info['chain1']=chain1
    info['chain2']=chain2
    info['len1']=len1
    info['len2']=len2
    info['class1']=class1
    info['class2']=class2
    
    return info

def get_pdb_chains(pdb):
    pdb_parser = Bio.PDB.PDBParser(QUIET = True)
    pdb_struct = pdb_parser.get_structure("reference", pdb)[0]
    chain=[]
    for c in pdb_struct:
        chain.append(c.id)
    return chain
#ATOM   3312  CA
#ATOM   3315  CB  ALA H 126     -21.802  31.674  73.597  1.00 58.05           C  

def make_two_chain_pdb(pdb,group1,group2): #renumber from 1
    pdb_parser = Bio.PDB.PDBParser(QUIET = True)
    pdb_struct = pdb_parser.get_structure("reference", pdb)[0]
    for c in pdb_struct:
        if c.id in group1:
            c.id='A'
        if c.id in group2:
            c.id='B'
    (code,outfile)=tempfile.mkstemp()
    io=Bio.PDB.PDBIO()
    io.set_structure(pdb_struct)
    io.save(outfile)
    exec_path=os.path.dirname(os.path.abspath(sys.argv[0]))    
    cmd=exec_path + '/scripts/renumber_pdb.pl ' + outfile
    os.system(cmd)
    os.remove(outfile)
    return outfile +'.renum'

def change_chain(pdb_string,chain):
    new_str=[];
    for line in pdb_string:
        s=list(line);
        s[21]=chain;
        new_str.append("".join(s));
    return "\n".join(new_str);

def make_two_chain_pdb_perm(pdb,group1,group2): #not yet ready
    pdb_chains={}
    f=open(pdb);
    for line in f.readlines():
        if line[0:4] == "ATOM":
            #       s=list(line);
            #print line
            chain=line[21]
            atom=line[13:16]
            resnum=int(line[22:26])
            #        print atom + ':' + str(resnum) +':'
            if chain not in pdb_chains:
                pdb_chains[chain]=[]
            pdb_chains[chain].append(line)
 #       print line
 #       s[21]='B'
 #       print "".join(s)
#        print chain


    f.close()
    #sys.exit()
    (code,outfile)=tempfile.mkstemp()
    f=open(outfile,'w')
    for c in group1:
     #   print pdb_chains[c]
        f.write(change_chain(pdb_chains[c],"A"))
    f.write("TER\n");
    for c in group2:
        f.write(change_chain(pdb_chains[c],"B"))
    f.close();
    #print outfile
    exec_path=os.path.dirname(os.path.abspath(sys.argv[0]))    
    cmd=exec_path + '/scripts/renumber_pdb.pl ' + outfile
    os.system(cmd)
    os.remove(outfile)
    return outfile +'.renum'
    
def main():

    parser=ArgumentParser(description="DockQ - Quality measure for protein-protein docking models")
    parser.add_argument('model',metavar='<model>',type=str,nargs=1,help='path to model file')
    parser.add_argument('native',metavar='<native>',type=str,nargs=1,help='path to native file')
    parser.add_argument('-capri_peptide',default=False,action='store_true',help='use version for capri_peptide (DockQ cannot not be trusted for this setting)')
    parser.add_argument('-short',default=False,action='store_true',help='short output')
    parser.add_argument('-verbose',default=False,action='store_true',help='talk a lot!')
    parser.add_argument('-quiet',default=False,action='store_true',help='keep quiet!')
    parser.add_argument('-useCA',default=False,action='store_true',help='use CA instead of backbone')
    parser.add_argument('-skip_check',default=False,action='store_true',help='skip initial check fo speed up on two chain examples')
    parser.add_argument('-no_needle',default=False,action='store_true',help='do not use global alignment to fix residue numbering between native and model during chain permutation (use only in case needle is not installed, and the residues between the chains are identical')
    parser.add_argument('-perm1',default=False,action='store_true',help='use all chain1 permutations to find maximum DockQ (number of comparisons is n! = 24, if combined with -perm2 there will be n!*m! combinations')
    parser.add_argument('-perm2',default=False,action='store_true',help='use all chain2 permutations to find maximum DockQ (number of comparisons is n! = 24, if combined with -perm1 there will be n!*m! combinations')
#    parser.add_argument('-comb',default=False,action='store_true',help='use all cyclicchain permutations to find maximum DockQ (number of comparisons is n!*m! = 24*24 = 576 for two tetramers interacting')
    parser.add_argument('-model_chain1',metavar='model_chain1', type=str,nargs='+', help='pdb chain order to group together partner 1')
    parser.add_argument('-model_chain2',metavar='model_chain2', type=str,nargs='+', help='pdb chain order to group together partner 2 (complement to partner 1 if undef)')
    parser.add_argument('-native_chain1',metavar='native_chain1', type=str,nargs='+', help='pdb chain order to group together from native partner 1')
    parser.add_argument('-native_chain2',metavar='native_chain2', type=str,nargs='+', help='pdb chain order to group together from native partner 2 (complement to partner 1 if undef)')


    args = parser.parse_args()
    #bio_ver=1.64
    bio_ver=1.61
    if(float(Bio.__version__) < bio_ver):
        print("Biopython version (%s) is too old need at least >=%.2f" % (Bio.__version__,bio_ver))
        sys.exit()

#    if(len(sys.argv)!=3):
#        print "Usage: ./Dock.py <model> <native>"
#        sys.exit()

#    print args
#    print args.model[0]
#    sys.exit()
#    model=sys.argv[1]
#    native=sys.argv[2]

    exec_path=os.path.dirname(os.path.abspath(sys.argv[0]))
    fix_numbering=exec_path + '/scripts/fix_numbering.pl'
    model=args.model[0]
    model_in=model
    native=args.native[0]
    native_in=native
    use_CA_only=args.useCA
    capri_peptide=args.capri_peptide

    model_chains=[]
    native_chains=[]
    best_info=''
    if(not args.skip_check):
        model_chains=get_pdb_chains(model)
        native_chains=get_pdb_chains(native)
    files_to_clean=[]

#    print native_chains
    if((len(model_chains) > 2 or len(native_chains) > 2) and
       (args.model_chain1 == None and args.native_chain1 == None)):
        print("Multi-chain model need sets of chains to group\nuse -native_chain1 and/or -model_chain1 if you want a different mapping than 1-1")
        print("Model chains  : " + str(model_chains))
        print("Native chains : " + str(native_chains))
        sys.exit()
    if not args.skip_check and (len(model_chains) < 2 or len(native_chains)< 2):
        print("Need at least two chains in the two inputs\n");
        sys.exit()
        
    if len(model_chains) > 2 or len(native_chains)> 2:
        group1=model_chains[0]
        group2=model_chains[1]
        nat_group1=native_chains[0]
        nat_group2=native_chains[1]
        if(args.model_chain1 != None):
            group1=args.model_chain1
            nat_group1=group1
            if(args.model_chain2 != None):
                group2=args.model_chain2
            else:
                #will use the complement from group1
                group2=[]
                for c in model_chains:
                    if c not in group1:
                        group2.append(c)
            nat_group1=group1
            nat_group2=group2
            

        if(args.native_chain1 != None):
            nat_group1=args.native_chain1
            if(args.native_chain2 != None):
                nat_group2=args.native_chain2
            else:
                #will use the complement from group1
                nat_group2=[]
                for c in native_chains:
                    if c not in nat_group1:
                        nat_group2.append(c)
                        
        if(args.model_chain1 == None):
            group1=nat_group1
            group2=nat_group2

        #print group1
        #print group2

        #print "native"
        #print nat_group1
        #print nat_group2
        if(args.verbose):
            print('Merging ' + ''.join(group1) + ' -> ' + ''.join(nat_group1) + ' to chain A')
            print('Merging ' + ''.join(group2) + ' -> ' + ''.join(nat_group2) + ' to chain B')
        native=make_two_chain_pdb_perm(native,nat_group1,nat_group2)
        files_to_clean.append(native)
        pe=0
        if args.perm1 or args.perm2:
            best_DockQ=-1;
            best_g1=[]
            best_g2=[]
            
            iter_perm1=itertools.combinations(group1,len(group1))
            iter_perm2=itertools.combinations(group2,len(group2))
            if args.perm1:
                iter_perm1=itertools.permutations(group1)
            if args.perm2:
                iter_perm2=itertools.permutations(group2)
               
            combos1=[];
            combos2=[];
            for g1 in iter_perm1:#_temp:
                combos1.append(g1)
            for g2 in iter_perm2:
                combos2.append(g2)

            for g1 in combos1:
                for g2 in combos2:
                    pe=pe+1
                   # print str(g1)+' '+str(g2)
#            print pe
#            print group1
#            print group2
            pe_tot=pe
            pe=1
            #sys.exit()
            if args.verbose:
                print('Starting chain order permutation search (number of permutations: ' + str(pe_tot) + ')')
            for g1 in combos1:
                for g2 in combos2:
                #g2=group2    
                    model_renum=make_two_chain_pdb_perm(model_in,g1,g2)
                    model_fixed=model_renum
                    if not args.no_needle:
                        fix_numbering_cmd=fix_numbering + ' ' + model_renum + ' ' + native + ' > /dev/null'
                        model_fixed=model_renum + ".fixed"
                        #                print fix_numbering_cmd
                        os.system(fix_numbering_cmd)
                        os.remove(model_renum)
                        if not os.path.exists(model_fixed):
                            print('If you are sure the residues are identical you can use the options -no_needle')
                            sys.exit()
                    test_dict=calc_DockQ(model_fixed,native,use_CA_only)
                    os.remove(model_fixed)
                    if not args.quiet:
                        print(str(pe)+'/'+str(pe_tot) + ' ' + ''.join(g1) + ' -> ' + ''.join(g2) + ' ' + str(test_dict['DockQ']))
                    if(test_dict['DockQ'] > best_DockQ):
                        best_DockQ=test_dict['DockQ'];
                        info=test_dict
                        best_g1=g1
                        best_g2=g2
                        best_info='Best score ( ' + str(best_DockQ) +' ) found for model -> native, chain1:' + ''.join(best_g1) + ' -> ' + ''.join(nat_group1) + ' chain2:' + ''.join(best_g2) + ' -> ' + ''.join(nat_group2)
                        
                        if args.verbose:
                            print(best_info)
                        if not args.quiet:    
                            print("Current best " + str(best_DockQ))
                    pe=pe+1
            if not args.quiet:
                print(best_info)        
#            print 'Best score ( ' + str(best_DockQ) +' ) found for ' + str(best_g1) + ' ' + str(best_g2)
        else:
            model_renum=make_two_chain_pdb_perm(model,group1,group2)
            model_fixed=model_renum
            if not args.no_needle:
                fix_numbering_cmd=fix_numbering + ' ' + model_renum + ' ' + native + ' > /dev/null'
                model_fixed=model_renum + ".fixed"
#                print fix_numbering_cmd
                os.system(fix_numbering_cmd)
                os.remove(model_renum)
                if not os.path.exists(model_fixed):
                    print('If you are sure the residues are identical you can use the options -no_needle')
                    sys.exit()
            info=calc_DockQ(model_fixed,native,use_CA_only)

            #os.system('cp ' + native + ' native_multichain.pdb')
            #os.system('cp ' + model_fixed + ' .')
            os.remove(model_fixed)
#            files_to_clean.append(model)
#            files_to_clean.append(model_fixed)
   #    sys.exit()
     
 #   print native
 #   print model
    else:
        info=calc_DockQ(model,native,use_CA_only=use_CA_only,capri_peptide=capri_peptide) #False):
#        info=calc_DockQ(model,native,use_CA_only=)
    
    irms=info['irms']
    Lrms=info['Lrms']
    fnat=info['fnat']
    DockQ=info['DockQ']
    fnonnat=info['fnonnat']
    
    if(args.short):
        if capri_peptide:
            print(("DockQ-capri_peptide %.3f Fnat %.3f iRMS %.3f LRMS %.3f Fnonnat %.3f %s %s %s" % (DockQ,fnat,irms,Lrms,fnonnat,model_in,native_in,best_info)))
        else:
            print(("DockQ %.3f Fnat %.3f iRMS %.3f LRMS %.3f Fnonnat %.3f %s %s %s" % (DockQ,fnat,irms,Lrms,fnonnat,model_in,native_in,best_info)))

    else:
        if capri_peptide:
            print('****************************************************************')
            print('*                DockQ-CAPRI peptide                           *')
            print('*   Do not trust any thing you read....                        *')
            print('*   OBS THE DEFINITION OF Fnat and iRMS are different for      *')
            print('*   peptides in CAPRI                                          *')
            print('*                                                              *')
            print('*   For the record:                                            *')
            print('*   Definition of contact <4A all heavy atoms (Fnat)           *')
            print('*   Definition of interface <8A CB (iRMS)                      *')
            print('*   For comments, please email: bjorn.wallner@.liu.se          *')
            print('****************************************************************')
        else:
            print('****************************************************************')
            print('*                       DockQ                                  *')
            print('*   Scoring function for protein-protein docking models        *')
            print('*   Statistics on CAPRI data:                                  *')
            print('*    0.00 <= DockQ <  0.23 - Incorrect                         *')
            print('*    0.23 <= DockQ <  0.49 - Acceptable quality                *')
            print('*    0.49 <= DockQ <  0.80 - Medium quality                    *')
            print('*            DockQ >= 0.80 - High quality                      *')  
            print('*   Reference: Sankar Basu and Bjorn Wallner, DockQ: A quality *')
            print('*   measure for protein-protein docking models, submitted      *')
            print('*                                                              *')
            print('*   For the record:                                            *')
            print('*   Definition of contact <5A (Fnat)                           *')
            print('*   Definition of interface <10A all heavy atoms (iRMS)        *')
            print('*   For comments, please email: bjorn.wallner@.liu.se          *')
            print('*                                                              *')
            print('****************************************************************')
        print(("Model  : %s" % model_in))
        print(("Native : %s" % native_in))
        if len(best_info):
            print(best_info)
        print('Number of equivalent residues in chain ' + info['chain1'] + ' ' + str(info['len1']) + ' (' + info['class1'] + ')')
        print('Number of equivalent residues in chain ' + info['chain2'] + ' ' + str(info['len2']) + ' (' + info['class2'] + ')')
        print(("Fnat %.3f %d correct of %d native contacts" % (info['fnat'],info['nat_correct'],info['nat_total'])))
        print(("Fnonnat %.3f %d non-native of %d model contacts" % (info['fnonnat'],info['nonnat_count'],info['model_total'])))
        print(("iRMS %.3f" % irms))
        print(("LRMS %.3f" % Lrms))
       # print 'CAPRI ' + capri_class(fnat,irms,Lrms,capri_peptide=capri_peptide)
        peptide_suffix=''
        if capri_peptide:
            peptide_suffix='_peptide'
        #print('CAPRI use DockQ instead.')
        #print(('CAPRI{} {}'.format(peptide_suffix,capri_class(fnat,irms,Lrms,capri_peptide=capri_peptide))))
        #print('DockQ_CAPRI ' + capri_class_DockQ(DockQ,capri_peptide=capri_peptide))
        peptide_disclaimer=''
        if capri_peptide:
            peptide_disclaimer='DockQ not reoptimized for CAPRI peptide evaluation'
        print(("DockQ {:.3f} {}".format(DockQ,peptide_disclaimer)))

    for f in files_to_clean:
        os.remove(f)

if __name__ == '__main__':
  main()    

AbDock/DockQ/DockQ.py
pengzhangzhi/ab_opt
from src.tools.runner.design_for_pdb import args_from_cmdline, dock_for_pdb

if __name__ == '__main__':
    dock_for_pdb(args_from_cmdline())

AbDock/dock_pdb.py
pengzhangzhi/ab_opt
from src.tools.runner.design_for_testset import main, parse_args, eval_all

if __name__ == "__main__":
    args = parse_args()
    if args.eval_all:
        eval_all(args)
    else:
        main(args)

AbDock/dock_testset.py
pengzhangzhi/ab_opt
import ray
from src.tools.eval.run import main, parse_args

if __name__ == '__main__':
    ray.init()
    args = parse_args()
    main(args)

AbDock/eval.py
pengzhangzhi/ab_opt
from easydict import EasyDict
import numpy as np
import os
import pandas as pd
import glob
import subprocess
import ray
import argparse

from src.tools.relax.run import main as relax_main
from src.tools.eval.run import main as eval_main
from src.tools.relax.run import pipeline_pyrosetta_fixbb,pipeline_openmm_pyrosetta

def seq_design_batch(save_dir, pose_dir, nums, heavy_chain_id, checkpoint_path, contig='',pdb_suffix='rosetta'):
    """
    run seq design procedure for sampled docking poses.
    """
    # initialize ray with the number of GPUs to use

    # specify the number of GPUs to use for each remote function call
    @ray.remote(num_gpus=1/args.process_per_gpu, num_cpus=1)
    def dock_pdb(path: str, save_dir: str, nums: int, heavy_chain_id: str):
        print(f"Processing pdb {path}")
        pdb_id = path.split('/')[-1].split('_')[0]
        cmd = f'python dock_pdb.py --pdb_path {path} -c configs/test/seq_design.yml\
            -ck {checkpoint_path} -o {save_dir} -n {nums} -b {nums} --heavy {heavy_chain_id} \
                --label_heavy_as_cdr'
        cmd += f' --contig {contig}' if contig else ''
        print(cmd)
        os.system(cmd)
        return path
    # create a list of remote function calls
    remote_calls = []
    for path in glob.glob(pose_dir+f'/*_{pdb_suffix}.pdb'):
        print(f"Submitting job for pdb {path}")
        remote_calls.append(dock_pdb.remote(path, save_dir, nums, heavy_chain_id))
    
    # retrieve the results from the remote function calls
    ray.get(remote_calls)
    summarize_seqs(save_dir)
        
def summarize_seqs(design_dir):
    out_df = pd.DataFrame(columns=['pdb_id', 'AAR'])
    
    for path in glob.glob(design_dir+'/*/aa.csv'):
        pdb_id = path.split('/')[-2]
        df = pd.read_csv(path)
        aar = df['AAR'].mean()
        out_df = out_df.append({'pdb_id': pdb_id, 'AAR': aar}, ignore_index=True) 
    out_df.to_csv(design_dir+'/summary.csv', index=False)
    return out_df

def gen_poses(native_path, out_dir, nums, checkpoint_path):
    cmd = f'python dock_pdb.py --pdb_path {native_path} -c configs/test/dock_design_single.yml\
                -ck {checkpoint_path} -o {out_dir} -n {nums} -b {nums}\
                   '
    print(cmd)
    os.system(cmd)

def dock_seqs(design_dir, out_dir, nums, heavy_chain_id, checkpoint_path, pdb_suffix):
    """
    Dock the designed sequences to the antigen.
    cmd:
    py dock_pdb.py --pdb_path wet_experiments/opt_ab/results/seq_design/7bsd/seq_design/0117_patch.pdb_/H_CDR3/0000_patch.pdb -c configs/test/dock_design_single.yml -ck reproduction/dock_single_cdr/250000.pt -o wet_experiments/opt_ab/results/screening -n 100 -b 100 --heavy A --label_heavy_as_cdr
    """
    @ray.remote(num_gpus=1/args.process_per_gpu, num_cpus=1)
    def dock_seq(path, id):
        cmd = f'python dock_pdb.py --id {id} --pdb_path {path} -c configs/test/dock_design_single.yml\
                -ck {checkpoint_path} -o {out_dir} -n {nums} -b {nums}\
                    --heavy {heavy_chain_id} --label_heavy_as_cdr'
        print(cmd)
        os.system(cmd)

    
    tasks = []
    paths = glob.glob(design_dir+f'/*_{pdb_suffix}.pdb_/H_CDR3/0000.pdb')
    if len(paths) == 0:
        raise ValueError(f'No pdb found in {design_dir}')
    for path in paths:
        id = path.split('/')[-3]
        print(id)
        tasks.append(dock_seq.remote(path, id))
        
    ray.get(tasks)
    ray.shutdown()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Script description')    
    
    parser.add_argument('--num_gpus', type=int, default=7, help='Number of GPUs')
    parser.add_argument('--process_per_gpu', type=int, default=1, help='Number of processes per GPU')
    parser.add_argument('--docked_pose_dir', type=str, default='wet_experiments/opt_ab/co_optimization/results/dock_single/7bsd_A_B_G.pdb_/H_CDR3', help='Docked pose directory')
    parser.add_argument('--seq_design_dir', type=str, default='wet_experiments/opt_ab/co_optimization/results/seq_design_fixed_pos/mutation/CDRH3_7_9', help='Sequence design directory')
    parser.add_argument('--design_model_ckpt', type=str, default='reproduction/seq_design_fixed_pos/300000.pt', help='Design model checkpoint')
    parser.add_argument('--design_contig', type=str, default='', help='Design contig')
    parser.add_argument('--screen_dir', type=str, default='wet_experiments/opt_ab/co_optimization/results/screening/seq_design_fixed_pos/mutation/CDRH3_7_9', help='Screening directory')
    parser.add_argument('--dock_model_ckpt', type=str, default='reproduction/dock_single_cdr/250000.pt', help='Dock model checkpoint')
    parser.add_argument('--heavy_chain_id', type=str, default='A', help='Heavy chain ID')
    parser.add_argument('--nums', type=int, default=100, help='Number of sequences')
    parser.add_argument('--pdb_suffix', type=str, default='rosetta', help='PDB suffix')

    args = parser.parse_args()
    os.makedirs(args.seq_design_dir, exist_ok=True)
    os.makedirs(args.screen_dir, exist_ok=True)
    
    ray.init(num_gpus=args.num_gpus)
    
    
    # run rosetta relax
    relax_main(EasyDict({
            'root': os.path.dirname(args.docked_pose_dir),
            "pipeline":pipeline_openmm_pyrosetta
    }))
    # run rosetta ddg estimation
    eval_main(args=EasyDict({
        'root': os.path.dirname(args.docked_pose_dir),
        'pfx': 'rosetta',
        'no_energy':False,
    }))

    
    seq_design_batch(
        save_dir=args.seq_design_dir,
        pose_dir=args.docked_pose_dir,
        nums=args.nums,
        heavy_chain_id=args.heavy_chain_id,
        checkpoint_path=args.design_model_ckpt,
        contig=args.design_contig,
        pdb_suffix=args.pdb_suffix,
    )

    dock_seqs(
        f'{args.seq_design_dir}/seq_design',
        args.screen_dir,
        nums=args.nums,
        heavy_chain_id=args.heavy_chain_id,
        checkpoint_path=args.dock_model_ckpt,
        pdb_suffix=args.pdb_suffix,
    )

    
AbDock/optimize_ab.py
pengzhangzhi/ab_opt
import ray
from src.tools.relax.run import run_relax, main, parse_args


if __name__ == '__main__':
    ray.init()
    args = parse_args()
    main(args)
AbDock/run_relax.py
pengzhangzhi/ab_opt
from .sabdab import SAbDabDataset
from .custom import CustomDataset

from ._base import get_dataset

AbDock/src/datasets/__init__.py
pengzhangzhi/ab_opt
from torch.utils.data import Dataset, ConcatDataset
from src.utils.transforms import get_transform


_DATASET_DICT = {}


def register_dataset(name):
    def decorator(cls):
        _DATASET_DICT[name] = cls
        return cls
    return decorator


def get_dataset(cfg):
    transform = get_transform(cfg.transform) if 'transform' in cfg else None
    return _DATASET_DICT[cfg.type](cfg, transform=transform)


@register_dataset('concat')
def get_concat_dataset(cfg):
    datasets = [get_dataset(d) for d in cfg.datasets]
    return ConcatDataset(datasets)


@register_dataset('balanced_concat')
class BalancedConcatDataset(Dataset):

    def __init__(self, cfg, transform=None):
        super().__init__()
        assert transform is None, 'transform is not supported.'
        self.datasets = [get_dataset(d) for d in cfg.datasets]
        self.max_size = max([len(d) for d in self.datasets])

    def __len__(self):
        return self.max_size * len(self.datasets)

    def __getitem__(self, idx):
        dataset_idx = idx // self.max_size
        return self.datasets[dataset_idx][idx % len(self.datasets[dataset_idx])]

AbDock/src/datasets/_base.py
pengzhangzhi/ab_opt
import os
import logging
import joblib
import pickle
import lmdb
from Bio import PDB
from Bio.PDB import PDBExceptions
from torch.utils.data import Dataset
from tqdm.auto import tqdm

from ..utils.protein import parsers
from .sabdab import _label_heavy_chain_cdr, _label_light_chain_cdr, _label_whole_chain_as_cdr
from ._base import register_dataset


def preprocess_antibody_structure(task,label_whole_heavy_chain_as_cdr=False):
    pdb_path = task['pdb_path']
    H_id = task.get('heavy_id', 'H')
    L_id = task.get('light_id', 'L')

    parser = PDB.PDBParser(QUIET=True)
    model = parser.get_structure(id, pdb_path)[0]

    all_chain_ids = [c.id for c in model]

    parsed = {
        'id': task['id'],
        'heavy': None,
        'heavy_seqmap': None,
        'light': None,
        'light_seqmap': None,
        'antigen': None,
        'antigen_seqmap': None,
    }
    try:
        if H_id in all_chain_ids:
            label_fn = _label_whole_chain_as_cdr if label_whole_heavy_chain_as_cdr else _label_heavy_chain_cdr
            (
                parsed['heavy'], 
                parsed['heavy_seqmap']
            ) = label_fn(*parsers.parse_biopython_structure(
                model[H_id],
                max_resseq = 113    # Chothia, end of Heavy chain Fv
            ))
        
        if L_id in all_chain_ids:
            (
                parsed['light'], 
                parsed['light_seqmap']
            ) = _label_light_chain_cdr(*parsers.parse_biopython_structure(
                model[L_id],
                max_resseq = 106    # Chothia, end of Light chain Fv
            ))

        if parsed['heavy'] is None and parsed['light'] is None:
            raise ValueError(
                f'Neither valid antibody H-chain or L-chain is found. '
                f'Please ensure that the chain id of heavy chain is "{H_id}" '
                f'and the id of the light chain is "{L_id}".'
            )

        
        ag_chain_ids = [cid for cid in all_chain_ids if cid not in (H_id, L_id)]
        if len(ag_chain_ids) > 0:
            chains = [model[c] for c in ag_chain_ids]
            (
                parsed['antigen'], 
                parsed['antigen_seqmap']
            ) = parsers.parse_biopython_structure(chains)

    except (
        PDBExceptions.PDBConstructionException, 
        parsers.ParsingException, 
        KeyError,
        ValueError,
    ) as e:
        logging.warning('[{}] {}: {}'.format(
            task['id'], 
            e.__class__.__name__, 
            str(e)
        ))
        return None

    return parsed


@register_dataset('custom')
class CustomDataset(Dataset):

    MAP_SIZE = 32*(1024*1024*1024)  # 32GB

    def __init__(self, structure_dir, transform=None, reset=False):
        super().__init__()
        self.structure_dir = structure_dir
        self.transform = transform

        self.db_conn = None
        self.db_ids = None
        self._load_structures(reset)

    @property
    def _cache_db_path(self):
        return os.path.join(self.structure_dir, 'structure_cache.lmdb')

    def _connect_db(self):
        self._close_db()
        self.db_conn = lmdb.open(
            self._cache_db_path,
            map_size=self.MAP_SIZE,
            create=False,
            subdir=False,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )
        with self.db_conn.begin() as txn:
            keys = [k.decode() for k in txn.cursor().iternext(values=False)]
            self.db_ids = keys

    def _close_db(self):
        if self.db_conn is not None:
            self.db_conn.close()
        self.db_conn = None
        self.db_ids = None
        
    def _load_structures(self, reset):
        all_pdbs = []
        for fname in os.listdir(self.structure_dir):
            if not fname.endswith('.pdb'): continue
            all_pdbs.append(fname)

        if reset or not os.path.exists(self._cache_db_path):
            todo_pdbs = all_pdbs
        else:
            self._connect_db()
            processed_pdbs = self.db_ids
            self._close_db()
            todo_pdbs = list(set(all_pdbs) - set(processed_pdbs))

        if len(todo_pdbs) > 0:
            self._preprocess_structures(todo_pdbs)
    
    def _preprocess_structures(self, pdb_list):
        tasks = []
        for pdb_fname in pdb_list:
            pdb_path = os.path.join(self.structure_dir, pdb_fname)
            tasks.append({
                'id': pdb_fname,
                'pdb_path': pdb_path,
            })

        data_list = joblib.Parallel(
            n_jobs = max(joblib.cpu_count() // 2, 1),
        )(
            joblib.delayed(preprocess_antibody_structure)(task)
            for task in tqdm(tasks, dynamic_ncols=True, desc='Preprocess')
        )

        db_conn = lmdb.open(
            self._cache_db_path,
            map_size = self.MAP_SIZE,
            create=True,
            subdir=False,
            readonly=False,
        )
        ids = []
        with db_conn.begin(write=True, buffers=True) as txn:
            for data in tqdm(data_list, dynamic_ncols=True, desc='Write to LMDB'):
                if data is None:
                    continue
                ids.append(data['id'])
                txn.put(data['id'].encode('utf-8'), pickle.dumps(data))

    def __len__(self):
        return len(self.db_ids)

    def __getitem__(self, index):
        self._connect_db()
        id = self.db_ids[index]
        with self.db_conn.begin() as txn:
            data = pickle.loads(txn.get(id.encode()))
        if self.transform is not None:
            data = self.transform(data)
        return data


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--dir', type=str, default='./data/custom')
    parser.add_argument('--reset', action='store_true', default=False)
    args = parser.parse_args()

    dataset = CustomDataset(
        structure_dir = args.dir,
        reset = args.reset,
    )
    print(dataset[0])
    print(len(dataset))
    
AbDock/src/datasets/custom.py
pengzhangzhi/ab_opt
import os
import random
import logging
import datetime
import pandas as pd
import joblib
import pickle
import lmdb
import subprocess
import torch
from Bio import PDB, SeqRecord, SeqIO, Seq
from Bio.PDB import PDBExceptions
from Bio.PDB import Polypeptide
from torch.utils.data import Dataset
from tqdm.auto import tqdm

from src.utils.transforms.patch import DISTVIOLATION

from ..utils.protein import parsers, constants
from ._base import register_dataset


ALLOWED_AG_TYPES = {
    'protein',
    'protein | protein',
    'protein | protein | protein',
    'protein | protein | protein | protein | protein',
    'protein | protein | protein | protein',
}

RESOLUTION_THRESHOLD = 4.0

TEST_ANTIGENS = [
    'sars-cov-2 receptor binding domain',
    'hiv-1 envelope glycoprotein gp160',
    'mers s',
    'influenza a virus',
    'cd27 antigen',
]
DOCK_TESTSETS = [
 '5mes_H_L_A',
 '1ic7_H_L_Y',
 '1ncb_H_L_N',
 '4lvn_C_B_A',
 '2adf_H_L_A',
 '4ydk_H_L_G',
 '1a2y_B_A_C',
 '3o2d_H_L_A',
 '1n8z_B_A_C',
 '4g6j_H_L_A',
 '1osp_H_L_O',
 '4fqj_H_L_A',
 '1uj3_B_A_C',
 '3uzq_A_a_B',
 '4h8w_H_L_G',
 '4dvr_H_L_G',
 '3s35_H_L_X',
 '3w9e_A_B_C',
 '5f9o_H_L_G',
 '1iqd_B_A_C',
 '4g6m_H_L_A',
 '3k2u_H_L_A',
 '4dtg_H_L_K',
 '5en2_A_B_C',
 '3mxw_H_L_A',
 '2dd8_H_L_S',
 '2ypv_H_L_A',
 '1a14_H_L_N',
 
 '4ffv_H_L_A',
 '3nid_E_F_C',
 '5ggs_C_D_Y',
 '3rkd_D_C_B',
 '4xnq_B_A_D',
 '2ghw_B_b_A',
 '2b2x_H_L_A',
 '4qci_H_L_C',
 '3l95_B_A_X',
 
 
 ## 
 '3hi6_H_L_A',
 '3cx5_U_V_P',
 '5bv7_C_B_A',
 '4ki5_E_F_M',
 '5nuz_H_L_D',
 '3bn9_D_C_B',
 '1fe8_H_L_A'
 ]
def nan_to_empty_string(val):
    if val != val or not val:
        return ''
    else:
        return val


def nan_to_none(val):
    if val != val or not val:
        return None
    else:
        return val


def split_sabdab_delimited_str(val):
    if not val:
        return []
    else:
        return [s.strip() for s in val.split('|')]


def parse_sabdab_resolution(val):
    if val == 'NOT' or not val or val != val:
        return None
    elif isinstance(val, str) and ',' in val:
        return float(val.split(',')[0].strip())
    else:
        return float(val)


def _aa_tensor_to_sequence(aa):
    return ''.join([Polypeptide.index_to_one(a.item()) for a in aa.flatten()])

def _label_whole_chain_as_cdr(data, seq_map, max_chain_length=120):
    """ label the whole chain to be CDRH3."""
    assert len(data['aa']) <= max_chain_length, f'Chain too long {len(data["aa"])}'
    if data is None or seq_map is None:
        return data, seq_map
    data['cdr_flag'] = torch.ones_like(data['aa']).fill_(constants.CDR.H3)
    return data, seq_map

def _label_heavy_chain_cdr(data, seq_map, max_cdr3_length=30):
    if data is None or seq_map is None:
        return data, seq_map

    # Add CDR labels
    cdr_flag = torch.zeros_like(data['aa'])
    for position, idx in seq_map.items():
        resseq = position[1]
        cdr_type = constants.ChothiaCDRRange.to_cdr('H', resseq)
        if cdr_type is not None:
            cdr_flag[idx] = cdr_type
    data['cdr_flag'] = cdr_flag

    # Add CDR sequence annotations
    data['H1_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.H1] )
    data['H2_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.H2] )
    data['H3_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.H3] )

    cdr3_length = (cdr_flag == constants.CDR.H3).sum().item()
    # Remove too long CDR3
    if cdr3_length > max_cdr3_length:
        cdr_flag[cdr_flag == constants.CDR.H3] = 0
        logging.warning(f'CDR-H3 too long {cdr3_length}. Removed.')
        return None, None

    # Filter: ensure CDR3 exists
    if cdr3_length == 0:
        logging.warning('No CDR-H3 found in the heavy chain.')
        return None, None

    return data, seq_map


def _label_light_chain_cdr(data, seq_map, max_cdr3_length=30):
    if data is None or seq_map is None:
        return data, seq_map
    cdr_flag = torch.zeros_like(data['aa'])
    for position, idx in seq_map.items():
        resseq = position[1]
        cdr_type = constants.ChothiaCDRRange.to_cdr('L', resseq)
        if cdr_type is not None:
            cdr_flag[idx] = cdr_type
    data['cdr_flag'] = cdr_flag

    data['L1_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.L1] )
    data['L2_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.L2] )
    data['L3_seq'] = _aa_tensor_to_sequence( data['aa'][cdr_flag == constants.CDR.L3] )

    cdr3_length = (cdr_flag == constants.CDR.L3).sum().item()
    # Remove too long CDR3
    if cdr3_length > max_cdr3_length:
        cdr_flag[cdr_flag == constants.CDR.L3] = 0
        logging.warning(f'CDR-L3 too long {cdr3_length}. Removed.')
        return None, None

    # Ensure CDR3 exists
    if cdr3_length == 0:
        logging.warning('No CDRs found in the light chain.')
        return None, None

    return data, seq_map


def preprocess_sabdab_structure(task):
    entry = task['entry']
    pdb_path = task['pdb_path']

    parser = PDB.PDBParser(QUIET=True)
    model = parser.get_structure(id, pdb_path)[0]

    parsed = {
        'id': entry['id'],
        'heavy': None,
        'heavy_seqmap': None,
        'light': None,
        'light_seqmap': None,
        'antigen': None,
        'antigen_seqmap': None,
    }
    try:
        if entry['H_chain'] is not None:
            (
                parsed['heavy'], 
                parsed['heavy_seqmap']
            ) = _label_heavy_chain_cdr(*parsers.parse_biopython_structure(
                model[entry['H_chain']],
                max_resseq = 113    # Chothia, end of Heavy chain Fv
            ))
        
        if entry['L_chain'] is not None:
            (
                parsed['light'], 
                parsed['light_seqmap']
            ) = _label_light_chain_cdr(*parsers.parse_biopython_structure(
                model[entry['L_chain']],
                max_resseq = 106    # Chothia, end of Light chain Fv
            ))

        if parsed['heavy'] is None and parsed['light'] is None:
            raise ValueError('Neither valid H-chain or L-chain is found.')
    
        if len(entry['ag_chains']) > 0:
            chains = [model[c] for c in entry['ag_chains']]
            (
                parsed['antigen'], 
                parsed['antigen_seqmap']
            ) = parsers.parse_biopython_structure(chains)

    except (
        PDBExceptions.PDBConstructionException, 
        parsers.ParsingException, 
        KeyError,
        ValueError,
    ) as e:
        logging.warning('[{}] {}: {}'.format(
            task['id'], 
            e.__class__.__name__, 
            str(e)
        ))
        return None

    return parsed


class SAbDabDataset(Dataset):

    MAP_SIZE = 32*(1024*1024*1024)  # 32GB

    def __init__(
        self, 
        summary_path = './data/sabdab_summary_all.tsv', 
        chothia_dir = './data/all_structures/chothia', 
        processed_dir = './data/processed',
        split = 'train',
        split_seed = 2022,
        transform = None,
        reset = False,
        must_have_antigen = False,
    ):
        super().__init__()
        self.must_have_antigen = must_have_antigen
        self.summary_path = summary_path
        self.chothia_dir = chothia_dir
        if not os.path.exists(chothia_dir):
            raise FileNotFoundError(
                f"SAbDab structures not found in {chothia_dir}. "
                "Please download them from http://opig.stats.ox.ac.uk/webapps/newsabdab/sabdab/archive/all/"
            )
        self.processed_dir = processed_dir
        os.makedirs(processed_dir, exist_ok=True)

        self.sabdab_entries = None
        self._load_sabdab_entries()

        self.db_conn = None
        self.db_ids = None
        self._load_structures(reset)

        self.clusters = None
        self.id_to_cluster = None
        self._load_clusters(reset)

        self.ids_in_split = None
        self._load_split(split, split_seed)

        self.transform = transform

    def _load_sabdab_entries(self):
        df = pd.read_csv(self.summary_path, sep='\t')
        entries_all = []
        for i, row in tqdm(
            df.iterrows(), 
            dynamic_ncols=True, 
            desc='Loading entries',
            total=len(df),
        ):
            entry_id = "{pdbcode}_{H}_{L}_{Ag}".format(
                pdbcode = row['pdb'],
                H = nan_to_empty_string(row['Hchain']),
                L = nan_to_empty_string(row['Lchain']),
                Ag = ''.join(split_sabdab_delimited_str(
                    nan_to_empty_string(row['antigen_chain'])
                ))
            )
            ag_chains = split_sabdab_delimited_str(
                nan_to_empty_string(row['antigen_chain'])
            )
            resolution = parse_sabdab_resolution(row['resolution'])
            entry = {
                'id': entry_id,
                'pdbcode': row['pdb'],
                'H_chain': nan_to_none(row['Hchain']),
                'L_chain': nan_to_none(row['Lchain']),
                'ag_chains': ag_chains,
                'ag_type': nan_to_none(row['antigen_type']),
                'ag_name': nan_to_none(row['antigen_name']),
                'date': datetime.datetime.strptime(row['date'], '%m/%d/%y'),
                'resolution': resolution,
                'method': row['method'],
                'scfv': row['scfv'],
            }

            # Filtering
            if (
                (entry['ag_type'] in ALLOWED_AG_TYPES or entry['ag_type'] is None)
                and (entry['resolution'] is not None and entry['resolution'] <= RESOLUTION_THRESHOLD)
            ):
                entries_all.append(entry)
        self.sabdab_entries = entries_all

    def _load_structures(self, reset):
        if not os.path.exists(self._structure_cache_path) or reset:
            if os.path.exists(self._structure_cache_path):
                os.unlink(self._structure_cache_path)
            self._preprocess_structures()

        with open(self._structure_cache_path + '-ids', 'rb') as f:
            self.db_ids = pickle.load(f)
        self.sabdab_entries = list(
            filter(
                lambda e: e['id'] in self.db_ids,
                self.sabdab_entries
            )
        )

    @property
    def _structure_cache_path(self):
        return os.path.join(self.processed_dir, 'structures.lmdb')
        
    def _preprocess_structures(self):
        tasks = []
        for entry in self.sabdab_entries:
            pdb_path = os.path.join(self.chothia_dir, '{}.pdb'.format(entry['pdbcode']))
            if not os.path.exists(pdb_path):
                logging.warning(f"PDB not found: {pdb_path}")
                continue
            tasks.append({
                'id': entry['id'],
                'entry': entry,
                'pdb_path': pdb_path,
            })

        data_list = joblib.Parallel(
            n_jobs = max(joblib.cpu_count() // 2, 1),
        )(
            joblib.delayed(preprocess_sabdab_structure)(task)
            for task in tqdm(tasks, dynamic_ncols=True, desc='Preprocess')
        )

        db_conn = lmdb.open(
            self._structure_cache_path,
            map_size = self.MAP_SIZE,
            create=True,
            subdir=False,
            readonly=False,
        )
        ids = []
        with db_conn.begin(write=True, buffers=True) as txn:
            for data in tqdm(data_list, dynamic_ncols=True, desc='Write to LMDB'):
                if data is None:
                    continue
                ids.append(data['id'])
                txn.put(data['id'].encode('utf-8'), pickle.dumps(data))

        with open(self._structure_cache_path + '-ids', 'wb') as f:
            pickle.dump(ids, f)

    @property
    def _cluster_path(self):
        return os.path.join(self.processed_dir, 'cluster_result_cluster.tsv')

    def _load_clusters(self, reset):
        if not os.path.exists(self._cluster_path) or reset:
            self._create_clusters()

        clusters, id_to_cluster = {}, {}
        with open(self._cluster_path, 'r') as f:
            for line in f.readlines():
                cluster_name, data_id = line.split()
                if cluster_name not in clusters:
                    clusters[cluster_name] = []
                clusters[cluster_name].append(data_id)
                id_to_cluster[data_id] = cluster_name
        self.clusters = clusters
        self.id_to_cluster = id_to_cluster

    def _create_clusters(self):
        cdr_records = []
        for id in self.db_ids:
            structure = self.get_structure(id)
            if structure['heavy'] is not None:
                cdr_records.append(SeqRecord.SeqRecord(
                    Seq.Seq(structure['heavy']['H3_seq']),
                    id = structure['id'],
                    name = '',
                    description = '',
                ))
            elif structure['light'] is not None:
                cdr_records.append(SeqRecord.SeqRecord(
                    Seq.Seq(structure['light']['L3_seq']),
                    id = structure['id'],
                    name = '',
                    description = '',
                ))
        fasta_path = os.path.join(self.processed_dir, 'cdr_sequences.fasta')
        SeqIO.write(cdr_records, fasta_path, 'fasta')

        cmd = ' '.join([
            'mmseqs', 'easy-cluster',
            os.path.realpath(fasta_path),
            'cluster_result', 'cluster_tmp',
            '--min-seq-id', '0.5',
            '-c', '0.8',
            '--cov-mode', '1',
        ])
        subprocess.run(cmd, cwd=self.processed_dir, shell=True, check=True)

    def _load_split(self, split, split_seed):
        assert split in ('train', 'val', 'test','dock_test')
        ids_test = [
            entry['id']
            for entry in self.sabdab_entries
            if entry['ag_name'] in TEST_ANTIGENS
        ]
        test_relevant_clusters = set([self.id_to_cluster[id] for id in ids_test])
        ids_docktest = [
            entry['id']
            for entry in self.sabdab_entries
            if entry['id'] in DOCK_TESTSETS
        ]

        dock_test_relevant_clusters = set([self.id_to_cluster[id] for id in ids_docktest])
        ids_train_val = [
            entry['id']
            for entry in self.sabdab_entries
            if (self.id_to_cluster[entry['id']] not in test_relevant_clusters) and \
                (self.id_to_cluster[entry['id']] not in dock_test_relevant_clusters)
        ]
        random.Random(split_seed).shuffle(ids_train_val)
        if split == 'test':
            self.ids_in_split = ids_test
        elif split == 'dock_test':
            self.ids_in_split = ids_docktest
            # print(self.ids_in_split)
        elif split == 'val':
            self.ids_in_split = ids_train_val[:20]
        else:
            self.ids_in_split = ids_train_val[20:]

    def _connect_db(self):
        if self.db_conn is not None:
            return
        self.db_conn = lmdb.open(
            self._structure_cache_path,
            map_size=self.MAP_SIZE,
            create=False,
            subdir=False,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )

    def get_structure(self, id):
        self._connect_db()
        with self.db_conn.begin() as txn:
            return pickle.loads(txn.get(id.encode()))

    def __len__(self):
        return len(self.ids_in_split)

    def __getitem__(self, index):
        id = self.ids_in_split[index]
        data = self.get_structure(id)
        if self.must_have_antigen and data.get('antigen', None) is None:
            return self.__getitem__((index + 1) % len(self))
        if self.transform is not None:
            try:
                data = self.transform(data)
            except DISTVIOLATION as e:
                return self.__getitem__((index + 1) % len(self))
        return data


@register_dataset('sabdab')
def get_sabdab_dataset(cfg, transform):
    return SAbDabDataset(
        summary_path = cfg.summary_path,
        chothia_dir = cfg.chothia_dir,
        processed_dir = cfg.processed_dir,
        split = cfg.split,
        split_seed = cfg.get('split_seed', 2022),
        transform = transform,
    )


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--split', type=str, default='train')
    parser.add_argument('--processed_dir', type=str, default='./data/processed')
    parser.add_argument('--reset', action='store_true', default=False)
    args = parser.parse_args()
    if args.reset:
        sure = input('Sure to reset? (y/n): ')
        if sure != 'y':
            exit()
    dataset = SAbDabDataset(
        processed_dir=args.processed_dir,
        split=args.split, 
        reset=args.reset
    )
    print(dataset[0])
    print(len(dataset), len(dataset.clusters))

AbDock/src/datasets/sabdab.py
pengzhangzhi/ab_opt
from .diffab import DiffusionAntibodyDesign

from ._base import get_model

AbDock/src/models/__init__.py
pengzhangzhi/ab_opt

_MODEL_DICT = {}


def register_model(name):
    def decorator(cls):
        _MODEL_DICT[name] = cls
        return cls
    return decorator


def get_model(cfg):
    return _MODEL_DICT[cfg.type](cfg)

AbDock/src/models/_base.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn

from src.modules.common.geometry import construct_3d_basis
from src.modules.common.so3 import rotation_to_so3vec
from src.modules.encoders.residue import ResidueEmbedding
from src.modules.encoders.pair import PairEmbedding
from src.modules.diffusion.dpm_full import FullDPM
from src.utils.protein.constants import max_num_heavyatoms, BBHeavyAtom, NUM_BB_ATOMS
from ._base import register_model


resolution_to_num_atoms = {
    'backbone+CB': 5,
    'full': max_num_heavyatoms
}


@register_model('diffab')
class DiffusionAntibodyDesign(nn.Module):

    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg

        num_atoms = resolution_to_num_atoms[cfg.get('resolution', 'full')]
        self.residue_embed = ResidueEmbedding(cfg.res_feat_dim, num_atoms)
        self.pair_embed = PairEmbedding(cfg.pair_feat_dim, num_atoms)

        self.diffusion = FullDPM(
            cfg.res_feat_dim,
            cfg.pair_feat_dim,
            **cfg.diffusion,
            num_bins=cfg.num_bins,
            dist_min=cfg.dist_min,
            dist_max=cfg.dist_max
        )

    def encode(self, batch, remove_structure, remove_sequence):
        """
        Returns:
            res_feat:   (N, L, res_feat_dim)
            pair_feat:  (N, L, L, pair_feat_dim)
        """
        # This is used throughout embedding and encoding layers
        #   to avoid data leakage.
        context_mask = torch.logical_and(
            batch['mask_heavyatom'][:, :, BBHeavyAtom.CA], 
            ~batch['generate_flag']     # Context means ``not generated''
        )

        structure_mask = context_mask if remove_structure else None
        sequence_mask = context_mask if remove_sequence else None

        res_feat = self.residue_embed(
            aa = batch['aa'],
            res_nb = batch['res_nb'],
            chain_nb = batch['chain_nb'],
            pos_atoms = batch['pos_heavyatom'],
            mask_atoms = batch['mask_heavyatom'],
            fragment_type = batch['fragment_type'],
            structure_mask = structure_mask,
            sequence_mask = sequence_mask,
        )

        pair_feat = self.pair_embed(
            aa = batch['aa'],
            res_nb = batch['res_nb'],
            chain_nb = batch['chain_nb'],
            pos_atoms = batch['pos_heavyatom'],
            mask_atoms = batch['mask_heavyatom'],
            structure_mask = structure_mask,
            sequence_mask = sequence_mask,
        )

        R = construct_3d_basis(
            batch['pos_heavyatom'][:, :, BBHeavyAtom.CA],
            batch['pos_heavyatom'][:, :, BBHeavyAtom.C],
            batch['pos_heavyatom'][:, :, BBHeavyAtom.N],
        )
        p = batch['pos_heavyatom'][:, :, BBHeavyAtom.CA]

        return res_feat, pair_feat, R, p
    
    def forward(self, batch):
        mask_generate = batch['generate_flag']
        # if self.cfg.get("bb_noise",False) \
        #     and self.cfg.get('train_sequence', True) is True:
        #     # add noise to backbone coordinate 
        #     bb_coords = batch['pos_heavyatom'][..., :NUM_BB_ATOMS,:]
        #     noise = self.cfg.bb_noise * torch.randn_like(bb_coords)
        #     batch['pos_heavyatom'][..., :NUM_BB_ATOMS,:] = bb_coords + noise
            
        if self.cfg.get('mask_ratio_min', False):
            random_mask = generate_random_mask_from(mask_generate, self.cfg.mask_ratio_min, self.cfg.mask_ratio_max)
            mask_generate = torch.logical_and(mask_generate, random_mask)
            batch['generate_flag'] = mask_generate
        mask_res = batch['mask']
        res_feat, pair_feat, R_0, p_0 = self.encode(
            batch,
            remove_structure = self.cfg.get('train_structure', True),
            remove_sequence = self.cfg.get('train_sequence', True)
        )
        v_0 = rotation_to_so3vec(R_0)
        s_0 = batch['aa']

        loss_dict = self.diffusion(
            v_0, p_0, s_0, res_feat, pair_feat, mask_generate, mask_res,
            denoise_structure = self.cfg.get('train_structure', True),
            denoise_sequence  = self.cfg.get('train_sequence', True),
        )
        return loss_dict

    @torch.no_grad()
    def sample(
        self, 
        batch, 
        sample_opt={
            'sample_structure': True,
            'sample_sequence': True,
            "contig": "",
        }
    ):
        mask_generate = batch['generate_flag']
        if sample_opt.get('sample_sequence', False)\
            and sample_opt['contig'] != '':
            mask = generate_mask_from_str(sample_opt['contig'], mask_generate)
            mask_generate = torch.logical_and(mask_generate, mask)
            batch['generate_flag'] = mask_generate
           
        mask_res = batch['mask']
        res_feat, pair_feat, R_0, p_0 = self.encode(
            batch,
            remove_structure = sample_opt.get('sample_structure', True),
            remove_sequence = sample_opt.get('sample_sequence', True)
        )
        v_0 = rotation_to_so3vec(R_0)
        s_0 = batch['aa']
        traj = self.diffusion.sample(v_0, p_0, s_0, res_feat, pair_feat, mask_generate, mask_res, **sample_opt)
        return traj

    @torch.no_grad()
    def optimize(
        self, 
        batch, 
        opt_step, 
        optimize_opt={
            'sample_structure': True,
            'sample_sequence': True,
        }
    ):
        mask_generate = batch['generate_flag']
        mask_res = batch['mask']
        res_feat, pair_feat, R_0, p_0 = self.encode(
            batch,
            remove_structure = optimize_opt.get('sample_structure', True),
            remove_sequence = optimize_opt.get('sample_sequence', True)
        )
        v_0 = rotation_to_so3vec(R_0)
        s_0 = batch['aa']

        traj = self.diffusion.optimize(v_0, p_0, s_0, opt_step, res_feat, pair_feat, mask_generate, mask_res, **optimize_opt)
        return traj


def generate_random_mask_from(tensor, mask_ratio_min, mask_ratio_max):
    """
    generate mask tensor with ratio between min and max.

    Args:
        mask_ratio_min (float): 
        mask_ratio_max (float): 
    
    """
    mask_ratio = float(torch.empty(1).uniform_(mask_ratio_min, mask_ratio_max))

    probs = torch.zeros_like(tensor.float()).fill_(mask_ratio)
    mask = torch.bernoulli(probs)
    mask = mask.bool()
    return mask



def generate_mask_from_str(str_input, tensor):
    """
    Generate a mask tensor from a string input. 
    The string input should be of the form 'start-end' where start and end are integers.
    
    Args:
        str_input (str): The string input, of the form 'start-end'
        tensor (torch.Tensor): template tensor to generate the mask from
    Returns:
        mask (torch.Tensor): The generated mask tensor, with the same shape as the input tensor. 
                                mask[..., start-1:end] = True, and other entries are False.
    """
    # Parse the string input to extract the start and end indices
    start, end = str_input.split('-')
    start, end = int(start), int(end)
    
    # Generate a mask tensor with all entries set to False
    mask = torch.zeros_like(tensor, dtype=torch.bool)
    
    # Set the entries from start to end to True
    mask[..., start-1:end] = True
    
    return mask
AbDock/src/models/diffab.py
pengzhangzhi/ab_opt
import torch
import torch.nn.functional as F

from src.utils.protein.constants import (
    BBHeavyAtom, 
    backbone_atom_coordinates_tensor,
    bb_oxygen_coordinate_tensor,
)
from .topology import get_terminus_flag


def safe_norm(x, dim=-1, keepdim=False, eps=1e-8, sqrt=True):
    out = torch.clamp(torch.sum(torch.square(x), dim=dim, keepdim=keepdim), min=eps)
    return torch.sqrt(out) if sqrt else out


def pairwise_distances(x, y=None, return_v=False):
    """
    Args:
        x:  (B, N, d)
        y:  (B, M, d)
    """
    if y is None: y = x
    v = x.unsqueeze(2) - y.unsqueeze(1)  # (B, N, M, d)
    d = safe_norm(v, dim=-1)
    if return_v:
        return d, v
    else:
        return d


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


def project_v2v(v, e, dim):
    """
    Description:
        Project vector `v` onto vector `e`.
    Args:
        v:  (N, L, 3).
        e:  (N, L, 3).
    """
    return (e * v).sum(dim=dim, keepdim=True) * e


def construct_3d_basis(center, p1, p2):
    """
    Args:
        center: (N, L, 3), usually the position of C_alpha.
        p1:     (N, L, 3), usually the position of C.
        p2:     (N, L, 3), usually the position of N.
    Returns
        A batch of orthogonal basis matrix, (N, L, 3, 3cols_index).
        The matrix is composed of 3 column vectors: [e1, e2, e3].
    """
    v1 = p1 - center    # (N, L, 3)
    e1 = normalize_vector(v1, dim=-1)

    v2 = p2 - center    # (N, L, 3)
    u2 = v2 - project_v2v(v2, e1, dim=-1)
    e2 = normalize_vector(u2, dim=-1)

    e3 = torch.cross(e1, e2, dim=-1)    # (N, L, 3)

    mat = torch.cat([
        e1.unsqueeze(-1), e2.unsqueeze(-1), e3.unsqueeze(-1)
    ], dim=-1)  # (N, L, 3, 3_index)
    return mat


def local_to_global(R, t, p):
    """
    Description:
        Convert local (internal) coordinates to global (external) coordinates q.
        q <- Rp + t
    Args:
        R:  (N, L, 3, 3).
        t:  (N, L, 3).
        p:  Local coordinates, (N, L, ..., 3).
    Returns:
        q:  Global coordinates, (N, L, ..., 3).
    """
    assert p.size(-1) == 3
    p_size = p.size()
    N, L = p_size[0], p_size[1]

    p = p.view(N, L, -1, 3).transpose(-1, -2)   # (N, L, *, 3) -> (N, L, 3, *)
    q = torch.matmul(R, p) + t.unsqueeze(-1)    # (N, L, 3, *)
    q = q.transpose(-1, -2).reshape(p_size)     # (N, L, 3, *) -> (N, L, *, 3) -> (N, L, ..., 3)
    return q


def global_to_local(R, t, q):
    """
    Description:
        Convert global (external) coordinates q to local (internal) coordinates p.
        p <- R^{T}(q - t)
    Args:
        R:  (N, L, 3, 3).
        t:  (N, L, 3).
        q:  Global coordinates, (N, L, ..., 3).
    Returns:
        p:  Local coordinates, (N, L, ..., 3).
    """
    assert q.size(-1) == 3
    q_size = q.size()
    N, L = q_size[0], q_size[1]

    q = q.reshape(N, L, -1, 3).transpose(-1, -2)   # (N, L, *, 3) -> (N, L, 3, *)
    p = torch.matmul(R.transpose(-1, -2), (q - t.unsqueeze(-1)))  # (N, L, 3, *)
    p = p.transpose(-1, -2).reshape(q_size)     # (N, L, 3, *) -> (N, L, *, 3) -> (N, L, ..., 3)
    return p


def apply_rotation_to_vector(R, p):
    return local_to_global(R, torch.zeros_like(p), p)


def compose_rotation_and_translation(R1, t1, R2, t2):
    """
    Args:
        R1,t1:  Frame basis and coordinate, (N, L, 3, 3), (N, L, 3).
        R2,t2:  Rotation and translation to be applied to (R1, t1), (N, L, 3, 3), (N, L, 3).
    Returns
        R_new <- R1R2
        t_new <- R1t2 + t1
    """
    R_new = torch.matmul(R1, R2)    # (N, L, 3, 3)
    t_new = torch.matmul(R1, t2.unsqueeze(-1)).squeeze(-1) + t1
    return R_new, t_new


def compose_chain(Ts):
    while len(Ts) >= 2:
        R1, t1 = Ts[-2]
        R2, t2 = Ts[-1]
        T_next = compose_rotation_and_translation(R1, t1, R2, t2)
        Ts = Ts[:-2] + [T_next]
    return Ts[0]


# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
def quaternion_to_rotation_matrix(quaternions):
    """
    Convert rotations given as quaternions to rotation matrices.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    quaternions = F.normalize(quaternions, dim=-1)
    r, i, j, k = torch.unbind(quaternions, -1)
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    return o.reshape(quaternions.shape[:-1] + (3, 3))


# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""
BSD License

For PyTorch3D software

Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

 * Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.

 * Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

 * Neither the name Meta nor the names of its contributors may be used to
   endorse or promote products derived from this software without specific
   prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
def quaternion_1ijk_to_rotation_matrix(q):
    """
    (1 + ai + bj + ck) -> R
    Args:
        q:  (..., 3)
    """
    b, c, d = torch.unbind(q, dim=-1)
    s = torch.sqrt(1 + b**2 + c**2 + d**2)
    a, b, c, d = 1/s, b/s, c/s, d/s

    o = torch.stack(
        (
            a**2 + b**2 - c**2 - d**2,  2*b*c - 2*a*d,  2*b*d + 2*a*c,
            2*b*c + 2*a*d,  a**2 - b**2 + c**2 - d**2,  2*c*d - 2*a*b,
            2*b*d - 2*a*c,  2*c*d + 2*a*b,  a**2 - b**2 - c**2 + d**2,
        ),
        -1,
    )
    return o.reshape(q.shape[:-1] + (3, 3))


def repr_6d_to_rotation_matrix(x):
    """
    Args:
        x:  6D representations, (..., 6).
    Returns:
        Rotation matrices, (..., 3, 3_index).
    """
    a1, a2 = x[..., 0:3], x[..., 3:6]
    b1 = normalize_vector(a1, dim=-1)
    b2 = normalize_vector(a2 - project_v2v(a2, b1, dim=-1), dim=-1)
    b3 = torch.cross(b1, b2, dim=-1)

    mat = torch.cat([
        b1.unsqueeze(-1), b2.unsqueeze(-1), b3.unsqueeze(-1)
    ], dim=-1)  # (N, L, 3, 3_index)
    return mat


def dihedral_from_four_points(p0, p1, p2, p3):
    """
    Args:
        p0-3:   (*, 3).
    Returns:
        Dihedral angles in radian, (*, ).
    """
    v0 = p2 - p1
    v1 = p0 - p1
    v2 = p3 - p2
    u1 = torch.cross(v0, v1, dim=-1)
    n1 = u1 / torch.linalg.norm(u1, dim=-1, keepdim=True)
    u2 = torch.cross(v0, v2, dim=-1)
    n2 = u2 / torch.linalg.norm(u2, dim=-1, keepdim=True)
    sgn = torch.sign( (torch.cross(v1, v2, dim=-1) * v0).sum(-1) )
    dihed = sgn*torch.acos( (n1 * n2).sum(-1).clamp(min=-0.999999, max=0.999999) )
    dihed = torch.nan_to_num(dihed)
    return dihed


def knn_gather(idx, value):
    """
    Args:
        idx:    (B, N, K)
        value:  (B, M, d)
    Returns:
        (B, N, K, d)
    """
    N, d = idx.size(1), value.size(-1)
    idx = idx.unsqueeze(-1).repeat(1, 1, 1, d)      # (B, N, K, d)
    value = value.unsqueeze(1).repeat(1, N, 1, 1)   # (B, N, M, d)
    return torch.gather(value, dim=2, index=idx)


def knn_points(q, p, K):
    """
    Args:
        q: (B, M, d)
        p: (B, N, d)
    Returns:
        (B, M, K), (B, M, K), (B, M, K, d)
    """
    _, L, _ = p.size()
    d = pairwise_distances(q, p)  # (B, N, M)
    dist, idx = d.topk(min(L, K), dim=-1, largest=False)  # (B, M, K), (B, M, K)
    return dist, idx, knn_gather(idx, p)


def angstrom_to_nm(x):
    return x / 10


def nm_to_angstrom(x):
    return x * 10


def get_backbone_dihedral_angles(pos_atoms, chain_nb, res_nb, mask):
    """
    Args:
        pos_atoms:  (N, L, A, 3).
        chain_nb:   (N, L).
        res_nb:     (N, L).
        mask:       (N, L).
    Returns:
        bb_dihedral:    Omega, Phi, and Psi angles in radian, (N, L, 3).
        mask_bb_dihed:  Masks of dihedral angles, (N, L, 3).
    """
    pos_N  = pos_atoms[:, :, BBHeavyAtom.N]   # (N, L, 3)
    pos_CA = pos_atoms[:, :, BBHeavyAtom.CA]
    pos_C  = pos_atoms[:, :, BBHeavyAtom.C]

    N_term_flag, C_term_flag = get_terminus_flag(chain_nb, res_nb, mask)  # (N, L)
    omega_mask = torch.logical_not(N_term_flag)
    phi_mask = torch.logical_not(N_term_flag)
    psi_mask = torch.logical_not(C_term_flag)

    # N-termini don't have omega and phi
    omega = F.pad(
        dihedral_from_four_points(pos_CA[:, :-1], pos_C[:, :-1], pos_N[:, 1:], pos_CA[:, 1:]), 
        pad=(1, 0), value=0,
    )
    phi = F.pad(
        dihedral_from_four_points(pos_C[:, :-1], pos_N[:, 1:], pos_CA[:, 1:], pos_C[:, 1:]),
        pad=(1, 0), value=0,
    )

    # C-termini don't have psi
    psi = F.pad(
        dihedral_from_four_points(pos_N[:, :-1], pos_CA[:, :-1], pos_C[:, :-1], pos_N[:, 1:]),
        pad=(0, 1), value=0,
    )

    mask_bb_dihed = torch.stack([omega_mask, phi_mask, psi_mask], dim=-1)
    bb_dihedral = torch.stack([omega, phi, psi], dim=-1) * mask_bb_dihed
    return bb_dihedral, mask_bb_dihed


def pairwise_dihedrals(pos_atoms):
    """
    Args:
        pos_atoms:  (N, L, A, 3).
    Returns:
        Inter-residue Phi and Psi angles, (N, L, L, 2).
    """
    N, L = pos_atoms.shape[:2]
    pos_N  = pos_atoms[:, :, BBHeavyAtom.N]   # (N, L, 3)
    pos_CA = pos_atoms[:, :, BBHeavyAtom.CA]
    pos_C  = pos_atoms[:, :, BBHeavyAtom.C]

    ir_phi = dihedral_from_four_points(
        pos_C[:,:,None].expand(N, L, L, 3), 
        pos_N[:,None,:].expand(N, L, L, 3), 
        pos_CA[:,None,:].expand(N, L, L, 3), 
        pos_C[:,None,:].expand(N, L, L, 3)
    )
    ir_psi = dihedral_from_four_points(
        pos_N[:,:,None].expand(N, L, L, 3), 
        pos_CA[:,:,None].expand(N, L, L, 3), 
        pos_C[:,:,None].expand(N, L, L, 3), 
        pos_N[:,None,:].expand(N, L, L, 3)
    )
    ir_dihed = torch.stack([ir_phi, ir_psi], dim=-1)
    return ir_dihed


def apply_rotation_matrix_to_rot6d(R, O):
    """
    Args:
        R:  (..., 3, 3)
        O:  (..., 6)
    Returns:
        Rotated 6D representation, (..., 6).
    """
    u1, u2 = O[..., :3, None], O[..., 3:, None] # (..., 3, 1)
    v1 = torch.matmul(R, u1).squeeze(-1)    # (..., 3)
    v2 = torch.matmul(R, u2).squeeze(-1)
    return torch.cat([v1, v2], dim=-1)


def normalize_rot6d(O):
    """
    Args:
        O:  (..., 6)
    """
    u1, u2 = O[..., :3], O[..., 3:]     # (..., 3)
    v1 = F.normalize(u1, p=2, dim=-1)   # (..., 3)
    v2 = F.normalize(u2 - project_v2v(u2, v1), p=2, dim=-1)
    return torch.cat([v1, v2], dim=-1)
    

def reconstruct_backbone(R, t, aa, chain_nb, res_nb, mask):
    """
    Args:
        R:  (N, L, 3, 3)
        t:  (N, L, 3)
        aa: (N, L)
        chain_nb:   (N, L)
        res_nb:     (N, L)
        mask:       (N, L)
    Returns:
        Reconstructed backbone atoms, (N, L, 4, 3).
    """
    N, L = aa.size()
    # atom_coords = restype_heavyatom_rigid_group_positions.clone().to(t) # (21, 14, 3)
    bb_coords = backbone_atom_coordinates_tensor.clone().to(t)  # (21, 3, 3)
    oxygen_coord = bb_oxygen_coordinate_tensor.clone().to(t)    # (21, 3)
    aa = aa.clamp(min=0, max=20)    # 20 for UNK

    bb_coords = bb_coords[aa.flatten()].reshape(N, L, -1, 3)    # (N, L, 3, 3)
    oxygen_coord = oxygen_coord[aa.flatten()].reshape(N, L, -1)  # (N, L, 3)
    bb_pos = local_to_global(R, t, bb_coords)   # Global coordinates of N, CA, C. (N, L, 3, 3).

    # Compute PSI angle
    bb_dihedral, _ = get_backbone_dihedral_angles(bb_pos, chain_nb, res_nb, mask)
    psi = bb_dihedral[..., 2]   # (N, L)
    # Make rotation matrix for PSI
    sin_psi = torch.sin(psi).reshape(N, L, 1, 1)
    cos_psi = torch.cos(psi).reshape(N, L, 1, 1)
    zero = torch.zeros_like(sin_psi)
    one = torch.ones_like(sin_psi)
    row1 = torch.cat([one, zero, zero], dim=-1)     # (N, L, 1, 3)
    row2 = torch.cat([zero, cos_psi, -sin_psi], dim=-1) # (N, L, 1, 3)
    row3 = torch.cat([zero, sin_psi, cos_psi], dim=-1)  # (N, L, 1, 3)
    R_psi = torch.cat([row1, row2, row3], dim=-2)       # (N, L, 3, 3)

    # Compute rotoation and translation of PSI frame, and position of O.
    R_psi, t_psi = compose_chain([
        (R, t), # Backbone
        (R_psi, torch.zeros_like(t)),       # PSI angle
    ])
    O_pos = local_to_global(R_psi, t_psi, oxygen_coord.reshape(N, L, 1, 3))

    bb_pos = torch.cat([bb_pos, O_pos], dim=2)  # (N, L, 4, 3)
    return bb_pos
    

def reconstruct_backbone_partially(pos_ctx, R_new, t_new, aa, chain_nb, res_nb, mask_atoms, mask_recons):
    """
    Args:
        pos:    (N, L, A, 3).
        R_new:  (N, L, 3, 3).
        t_new:  (N, L, 3).
        mask_atoms: (N, L, A).
        mask_recons:(N, L).
    Returns:
        pos_new:    (N, L, A, 3).
        mask_new:   (N, L, A).
    """
    N, L, A = mask_atoms.size()

    mask_res = mask_atoms[:, :, BBHeavyAtom.CA]
    pos_recons = reconstruct_backbone(R_new, t_new, aa, chain_nb, res_nb, mask_res) # (N, L, 4, 3)
    pos_recons = F.pad(pos_recons, pad=(0, 0, 0, A-4), value=0) # (N, L, A, 3)

    pos_new = torch.where(
        mask_recons[:, :, None, None].expand_as(pos_ctx),
        pos_recons, pos_ctx
    )   # (N, L, A, 3)

    mask_bb_atoms = torch.zeros_like(mask_atoms)
    mask_bb_atoms[:, :, :4] = True
    mask_new = torch.where(
        mask_recons[:, :, None].expand_as(mask_atoms),
        mask_bb_atoms, mask_atoms
    )

    return pos_new, mask_new


AbDock/src/modules/common/geometry.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F


def mask_zero(mask, value):
    return torch.where(mask, value, torch.zeros_like(value))


def clampped_one_hot(x, num_classes):
    mask = (x >= 0) & (x < num_classes) # (N, L)
    x = x.clamp(min=0, max=num_classes-1)
    y = F.one_hot(x, num_classes) * mask[...,None]  # (N, L, C)
    return y


class DistanceToBins(nn.Module):

    def __init__(self, dist_min=0.0, dist_max=20.0, num_bins=64, use_onehot=False):
        super().__init__()
        self.dist_min = dist_min
        self.dist_max = dist_max
        self.num_bins = num_bins
        self.use_onehot = use_onehot

        if use_onehot:
            offset = torch.linspace(dist_min, dist_max, self.num_bins)
        else:
            offset = torch.linspace(dist_min, dist_max, self.num_bins-1)    # 1 overflow flag
            self.coeff = -0.5 / ((offset[1] - offset[0]) * 0.2).item() ** 2  # `*0.2`: makes it not too blurred
        self.register_buffer('offset', offset)

    @property
    def out_channels(self):
        return self.num_bins 

    def forward(self, dist, dim, normalize=True):
        """
        Args:
            dist:   (N, *, 1, *)
        Returns:
            (N, *, num_bins, *)
        """
        assert dist.size()[dim] == 1
        offset_shape = [1] * len(dist.size())
        offset_shape[dim] = -1

        if self.use_onehot:
            diff = torch.abs(dist - self.offset.view(*offset_shape))  # (N, *, num_bins, *)
            bin_idx = torch.argmin(diff, dim=dim, keepdim=True)  # (N, *, 1, *)
            y = torch.zeros_like(diff).scatter_(dim=dim, index=bin_idx, value=1.0)
        else:
            overflow_symb = (dist >= self.dist_max).float()  # (N, *, 1, *)
            y = dist - self.offset.view(*offset_shape)  # (N, *, num_bins-1, *)
            y = torch.exp(self.coeff * torch.pow(y, 2))  # (N, *, num_bins-1, *)
            y = torch.cat([y, overflow_symb], dim=dim)  # (N, *, num_bins, *)
            if normalize:
                y = y / y.sum(dim=dim, keepdim=True)

        return y


class PositionalEncoding(nn.Module):
    
    def __init__(self, num_funcs=6):
        super().__init__()
        self.num_funcs = num_funcs
        self.register_buffer('freq_bands', 2.0 ** torch.linspace(0.0, num_funcs-1, num_funcs))
    
    def get_out_dim(self, in_dim):
        return in_dim * (2 * self.num_funcs + 1)

    def forward(self, x):
        """
        Args:
            x:  (..., d).
        """
        shape = list(x.shape[:-1]) + [-1]
        x = x.unsqueeze(-1) # (..., d, 1)
        code = torch.cat([x, torch.sin(x * self.freq_bands), torch.cos(x * self.freq_bands)], dim=-1)   # (..., d, 2f+1)
        code = code.reshape(shape)
        return code


class AngularEncoding(nn.Module):

    def __init__(self, num_funcs=3):
        super().__init__()
        self.num_funcs = num_funcs
        self.register_buffer('freq_bands', torch.FloatTensor(
            [i+1 for i in range(num_funcs)] + [1./(i+1) for i in range(num_funcs)]
        ))

    def get_out_dim(self, in_dim):
        return in_dim * (1 + 2*2*self.num_funcs)

    def forward(self, x):
        """
        Args:
            x:  (..., d).
        """
        shape = list(x.shape[:-1]) + [-1]
        x = x.unsqueeze(-1) # (..., d, 1)
        code = torch.cat([x, torch.sin(x * self.freq_bands), torch.cos(x * self.freq_bands)], dim=-1)   # (..., d, 2f+1)
        code = code.reshape(shape)
        return code


class LayerNorm(nn.Module):

    def __init__(self,
                 normal_shape,
                 gamma=True,
                 beta=True,
                 epsilon=1e-10):
        """Layer normalization layer
        See: [Layer Normalization](https://arxiv.org/pdf/1607.06450.pdf)
        :param normal_shape: The shape of the input tensor or the last dimension of the input tensor.
        :param gamma: Add a scale parameter if it is True.
        :param beta: Add an offset parameter if it is True.
        :param epsilon: Epsilon for calculating variance.
        """
        super().__init__()
        if isinstance(normal_shape, int):
            normal_shape = (normal_shape,)
        else:
            normal_shape = (normal_shape[-1],)
        self.normal_shape = torch.Size(normal_shape)
        self.epsilon = epsilon
        if gamma:
            self.gamma = nn.Parameter(torch.Tensor(*normal_shape))
        else:
            self.register_parameter('gamma', None)
        if beta:
            self.beta = nn.Parameter(torch.Tensor(*normal_shape))
        else:
            self.register_parameter('beta', None)
        self.reset_parameters()

    def reset_parameters(self):
        if self.gamma is not None:
            self.gamma.data.fill_(1)
        if self.beta is not None:
            self.beta.data.zero_()

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
        std = (var + self.epsilon).sqrt()
        y = (x - mean) / std
        if self.gamma is not None:
            y *= self.gamma
        if self.beta is not None:
            y += self.beta
        return y

    def extra_repr(self):
        return 'normal_shape={}, gamma={}, beta={}, epsilon={}'.format(
            self.normal_shape, self.gamma is not None, self.beta is not None, self.epsilon,
        )

AbDock/src/modules/common/layers.py
pengzhangzhi/ab_opt
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from functools import partial
import math
from typing import Optional, Callable, List, Tuple, Sequence
import numpy as np
import torch
import torch.nn as nn
from scipy.stats import truncnorm

from src.modules.common.layers import LayerNorm



DEFAULT_LMA_Q_CHUNK_SIZE=1024
DEFAULT_LMA_KV_CHUNK_SIZE=4096


def _prod(nums):
    out = 1
    for n in nums:
        out = out * n
    return out


def _calculate_fan(linear_weight_shape, fan="fan_in"):
    fan_out, fan_in = linear_weight_shape

    if fan == "fan_in":
        f = fan_in
    elif fan == "fan_out":
        f = fan_out
    elif fan == "fan_avg":
        f = (fan_in + fan_out) / 2
    else:
        raise ValueError("Invalid fan option")

    return f


def trunc_normal_init_(weights, scale=1.0, fan="fan_in"):
    shape = weights.shape
    f = _calculate_fan(shape, fan)
    scale = scale / max(1, f)
    a = -2
    b = 2
    std = math.sqrt(scale) / truncnorm.std(a=a, b=b, loc=0, scale=1)
    size = _prod(shape)
    samples = truncnorm.rvs(a=a, b=b, loc=0, scale=std, size=size)
    samples = np.reshape(samples, shape)
    with torch.no_grad():
        weights.copy_(torch.tensor(samples, device=weights.device))


def lecun_normal_init_(weights):
    trunc_normal_init_(weights, scale=1.0)


def he_normal_init_(weights):
    trunc_normal_init_(weights, scale=2.0)


def glorot_uniform_init_(weights):
    nn.init.xavier_uniform_(weights, gain=1)


def final_init_(weights):
    with torch.no_grad():
        weights.fill_(0.0)


def gating_init_(weights):
    with torch.no_grad():
        weights.fill_(0.0)


def normal_init_(weights):
    torch.nn.init.kaiming_normal_(weights, nonlinearity="linear")


def ipa_point_weights_init_(weights):
    with torch.no_grad():
        softplus_inverse_1 = 0.541324854612918
        weights.fill_(softplus_inverse_1)


class Linear(nn.Linear):
    """
    A Linear layer with built-in nonstandard initializations. Called just
    like torch.nn.Linear.
    Implements the initializers in 1.11.4, plus some additional ones found
    in the code.
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        bias: bool = True,
        init: str = "default",
        init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None,
    ):
        """
        Args:
            in_dim:
                The final dimension of inputs to the layer
            out_dim:
                The final dimension of layer outputs
            bias:
                Whether to learn an additive bias. True by default
            init:
                The initializer to use. Choose from:
                "default": LeCun fan-in truncated normal initialization
                "relu": He initialization w/ truncated normal distribution
                "glorot": Fan-average Glorot uniform initialization
                "gating": Weights=0, Bias=1
                "normal": Normal initialization with std=1/sqrt(fan_in)
                "final": Weights=0, Bias=0
                Overridden by init_fn if the latter is not None.
            init_fn:
                A custom initializer taking weight and bias as inputs.
                Overrides init if not None.
        """
        super(Linear, self).__init__(in_dim, out_dim, bias=bias)

        if bias:
            with torch.no_grad():
                self.bias.fill_(0)

        with torch.no_grad():
            if init_fn is not None:
                init_fn(self.weight, self.bias)
            else:
                if init == "default":
                    lecun_normal_init_(self.weight)
                elif init == "relu":
                    he_normal_init_(self.weight)
                elif init == "glorot":
                    glorot_uniform_init_(self.weight)
                elif init == "gating":
                    gating_init_(self.weight)
                    if bias:
                        self.bias.fill_(1.0)
                elif init == "normal":
                    normal_init_(self.weight)
                elif init == "final":
                    final_init_(self.weight)
                else:
                    raise ValueError("Invalid init string.")


class PerResiduePredictor(nn.Module):
    def __init__(self, no_bins, c_in, c_hidden):
        super(PerResiduePredictor, self).__init__()

        self.no_bins = no_bins
        self.c_in = c_in
        self.c_hidden = c_hidden

        self.layer_norm = LayerNorm(self.c_in)

        self.linear_1 = Linear(self.c_in, self.c_hidden, init="relu")
        self.linear_2 = Linear(self.c_hidden, self.c_hidden, init="relu")
        self.linear_3 = Linear(self.c_hidden, self.no_bins, init="final")

        self.relu = nn.ReLU()

    def forward(self, s):
        s = self.layer_norm(s)
        s = self.linear_1(s)
        s = self.relu(s)
        s = self.linear_2(s)
        s = self.relu(s)
        s = self.linear_3(s)

        return s

AbDock/src/modules/common/nn.py
pengzhangzhi/ab_opt
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn
from src.modules.common.nn import PerResiduePredictor
from src.modules.common.prmsd import softmax_cross_entropy

from src.modules.common.tensor_utils import permute_final_dims

class PerResidueLDDTCaPredictor(PerResiduePredictor):
    def __init__(self, no_bins, c_in, c_hidden):
        super().__init__(no_bins, c_in, c_hidden)
        
def compute_plddt(logits: torch.Tensor) -> torch.Tensor:
    num_bins = logits.shape[-1]
    bin_width = 1.0 / num_bins
    bounds = torch.arange(
        start=0.5 * bin_width, end=1.0, step=bin_width, device=logits.device
    )
    probs = torch.nn.functional.softmax(logits, dim=-1)
    pred_lddt_ca = torch.sum(
        probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape),
        dim=-1,
    )
    return pred_lddt_ca * 100


def lddt(
    all_atom_pred_pos: torch.Tensor,
    all_atom_positions: torch.Tensor,
    all_atom_mask: torch.Tensor,
    cutoff: float = 15.0,
    eps: float = 1e-10,
    per_residue: bool = True,
) -> torch.Tensor:
    n = all_atom_mask.shape[-2]
    dmat_true = torch.sqrt(
        eps
        + torch.sum(
            (
                all_atom_positions[..., None, :]
                - all_atom_positions[..., None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    dmat_pred = torch.sqrt(
        eps
        + torch.sum(
            (
                all_atom_pred_pos[..., None, :]
                - all_atom_pred_pos[..., None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )
    dists_to_score = (
        (dmat_true < cutoff)
        * all_atom_mask
        * permute_final_dims(all_atom_mask, (1, 0))
        * (1.0 - torch.eye(n, device=all_atom_mask.device))
    )

    dist_l1 = torch.abs(dmat_true - dmat_pred)

    score = (
        (dist_l1 < 0.5).type(dist_l1.dtype)
        + (dist_l1 < 1.0).type(dist_l1.dtype)
        + (dist_l1 < 2.0).type(dist_l1.dtype)
        + (dist_l1 < 4.0).type(dist_l1.dtype)
    )
    score = score * 0.25

    dims = (-1,) if per_residue else (-2, -1)
    norm = 1.0 / (eps + torch.sum(dists_to_score, dim=dims))
    score = norm * (eps + torch.sum(dists_to_score * score, dim=dims))

    return score



def lddt_loss(
    logits: torch.Tensor,
    all_atom_pred_pos: torch.Tensor,
    all_atom_positions: torch.Tensor,
    all_atom_mask: torch.Tensor,
    cutoff: float = 15.0,
    no_bins: int = 50,
    eps: float = 1e-10,
    **kwargs,
) -> torch.Tensor:
    """
    Compute the LDDT loss.

    Args:
        logits (torch.Tensor): A tensor of shape [B, L, no_bins],
        all_atom_pred_pos (torch.Tensor): A tensor of shape [B, L, 3], 
        all_atom_positions (torch.Tensor): A tensor of shape [B, L, 3], 
        all_atom_mask (torch.Tensor): A binary mask of shape [B, L, 1], where each element is either 0 or 1, 
            indicating whether the corresponding atom in the sequence should be included in the calculation 
            of the LDDT loss.
        cutoff (float, optional): The distance cutoff (in Angstroms) beyond which pairwise distances are 
            considered to be meaningless and ignored in the calculation of the LDDT loss. Defaults to 15.0.
        no_bins (int, optional): 

    Returns:
        torch.Tensor: _description_
    """
    score = lddt(
        all_atom_pred_pos, 
        all_atom_positions, 
        all_atom_mask, 
        cutoff=cutoff, 
        eps=eps
    )

    score = score.detach()

    bin_index = torch.floor(score * no_bins).long()
    bin_index = torch.clamp(bin_index, max=(no_bins - 1))
    lddt_ca_one_hot = torch.nn.functional.one_hot(
        bin_index, num_classes=no_bins
    )

    errors = softmax_cross_entropy(logits, lddt_ca_one_hot)
    all_atom_mask = all_atom_mask.squeeze(-1)
    loss = torch.sum(errors * all_atom_mask, dim=-1) / (
        eps + torch.sum(all_atom_mask, dim=-1)
    )


    # Average over the batch dimension
    loss = torch.mean(loss)

    return loss
AbDock/src/modules/common/plddt.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
from typing import Optional, Callable, List, Tuple, Sequence

from src.modules.common.layers import DistanceToBins, LayerNorm
from src.modules.common.nn import Linear, PerResiduePredictor
    
class PerResidueRMSDCaPredictor(PerResiduePredictor):
    def __init__(self, no_bins, c_in, c_hidden):
        super().__init__(no_bins, c_in, c_hidden)

def softmax_cross_entropy(logits, labels):
    loss = -1 * torch.sum(
        labels * torch.nn.functional.log_softmax(logits, dim=-1),
        dim=-1,
    )
    return loss

class pRMSDCa(nn.Module):
    def __init__(self, num_bins=20, dist_min=0.5, dist_max=19.5):
        """
        [dist_min, ..., dist_max]
        """
        super(pRMSDCa, self).__init__()

        self.num_bins = num_bins
        self.dist_min = dist_min
        self.dist_max = dist_max
        self.tobin = DistanceToBins(dist_min=self.dist_min, dist_max=self.dist_max, num_bins=self.num_bins, use_onehot=True)
        
    def compute_prmsd(self, logits: torch.Tensor) -> torch.Tensor:
        """
        Args:
            logits (torch.Tensor): [..., num_bins]

        Returns:
            torch.Tensor: [...,]
        """
        num_bins = logits.shape[-1]
        assert num_bins == self.num_bins
        bounds = torch.linspace(self.dist_min, self.dist_max, self.num_bins, device=logits.device)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        pred_rmsd_ca = torch.sum(
            probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape),
            dim=-1,
        )
        return pred_rmsd_ca

    def forward(self,prmsd_logits: torch.Tensor, rmsd: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        loss = self.calc_prmsd_loss(prmsd_logits, rmsd, mask)
        return loss
    
    def calc_prmsd_loss(self, prmsd_logits: torch.Tensor, rmsd: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            prmsd_logits (torch.Tensor): [B, ..., num_logits]
            rmsd (torch.Tensor): [B, ...]
            mask (torch.Tensor): [B, ...]
        
        Returns:
            torch.Tensor: 
        """
        eps: float = 1e-10

        
        rmsd_onehot = self.tobin(rmsd.unsqueeze(-1),dim=-1)
        errors = softmax_cross_entropy(prmsd_logits, rmsd_onehot) # [B, ...]
        loss = (errors * mask).sum() / (mask.sum() + eps)
        return loss

    @staticmethod
    def calc_per_rmsd(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pred (torch.Tensor): [B, L, 3]
            target (torch.Tensor): [B, L, 3]

        Returns:
            torch.Tensor: [B, L]
        """
        diff = pred - target
        diff = diff ** 2
        diff = torch.sum(diff, dim=-1)
        diff = torch.sqrt(diff)
        return diff

    @staticmethod
    def calc_rmsd(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pred (torch.Tensor): [B, L, 3]
            target (torch.Tensor): [B, L, 3]
            mask (torch.Tensor): [B, L]

        Returns:
            torch.Tensor: [B]
        """
        # Apply mask to predicted and target tensors
        pred_masked = pred * mask.unsqueeze(-1)
        target_masked = target * mask.unsqueeze(-1)

        # Calculate squared distance between masked predicted and target tensors
        squared_diff = torch.sum((pred_masked - target_masked)**2, dim=-1)

        # Count number of unmasked positions
        num_unmasked = torch.sum(mask, dim=-1)

        # Calculate RMSD for each example
        rmsd = torch.sqrt(torch.sum(squared_diff, dim=-1) / num_unmasked)

        return rmsd


AbDock/src/modules/common/prmsd.py
pengzhangzhi/ab_opt
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .geometry import quaternion_to_rotation_matrix


def log_rotation(R):
    trace = R[..., range(3), range(3)].sum(-1)
    if torch.is_grad_enabled():
        # The derivative of acos at -1.0 is -inf, so to stablize the gradient, we use -0.9999
        min_cos = -0.999
    else:
        min_cos = -1.0
    cos_theta = ( (trace-1) / 2 ).clamp_min(min=min_cos)
    sin_theta = torch.sqrt(1 - cos_theta**2)
    theta = torch.acos(cos_theta)
    coef = ((theta+1e-8)/(2*sin_theta+2e-8))[..., None, None]
    logR = coef * (R - R.transpose(-1, -2))
    return logR


def skewsym_to_so3vec(S):
    x = S[..., 1, 2]
    y = S[..., 2, 0]
    z = S[..., 0, 1]
    w = torch.stack([x,y,z], dim=-1)
    return w


def so3vec_to_skewsym(w):
    x, y, z = torch.unbind(w, dim=-1)
    o = torch.zeros_like(x)
    S = torch.stack([
        o, z, -y,
        -z, o, x,
        y, -x, o,
    ], dim=-1).reshape(w.shape[:-1] + (3, 3))
    return S


def exp_skewsym(S):
    x = torch.linalg.norm(skewsym_to_so3vec(S), dim=-1)
    I = torch.eye(3).to(S).view([1 for _ in range(S.dim()-2)] + [3, 3])
    
    sinx, cosx = torch.sin(x), torch.cos(x)
    b = (sinx + 1e-8) / (x + 1e-8)
    c = (1-cosx + 1e-8) / (x**2 + 2e-8)  # lim_{x->0} (1-cosx)/(x^2) = 0.5

    S2 = S @ S
    return I + b[..., None, None]*S + c[..., None, None]*S2


def so3vec_to_rotation(w):
    return exp_skewsym(so3vec_to_skewsym(w))


def rotation_to_so3vec(R):
    logR = log_rotation(R)
    w = skewsym_to_so3vec(logR)
    return w


def random_uniform_so3(size, device='cpu'):
    q = F.normalize(torch.randn(list(size)+[4,], device=device), dim=-1)    # (..., 4)
    return rotation_to_so3vec(quaternion_to_rotation_matrix(q))


class ApproxAngularDistribution(nn.Module):

    def __init__(self, stddevs, std_threshold=0.1, num_bins=8192, num_iters=1024):
        super().__init__()
        self.std_threshold = std_threshold
        self.num_bins = num_bins
        self.num_iters = num_iters
        self.register_buffer('stddevs', torch.FloatTensor(stddevs))
        self.register_buffer('approx_flag', self.stddevs <= std_threshold)
        self._precompute_histograms()

    @staticmethod
    def _pdf(x, e, L):
        """
        Args:
            x:  (N, )
            e:  Float
            L:  Integer
        """
        x = x[:, None]  # (N, *)
        c = ((1 - torch.cos(x)) / math.pi)  # (N, *)
        l = torch.arange(0, L)[None, :]  # (*, L)
        a = (2*l+1) * torch.exp(-l*(l+1)*(e**2))  # (*, L)
        b = (torch.sin( (l+0.5)* x ) + 1e-6) / (torch.sin( x / 2 ) + 1e-6) # (N, L)
        
        f = (c * a * b).sum(dim=1)
        return f

    def _precompute_histograms(self):
        X, Y = [], []
        for std in self.stddevs:
            std = std.item()
            x = torch.linspace(0, math.pi, self.num_bins)   # (n_bins,)
            y = self._pdf(x, std, self.num_iters)    # (n_bins,)
            y = torch.nan_to_num(y).clamp_min(0)
            X.append(x)
            Y.append(y)
        self.register_buffer('X', torch.stack(X, dim=0))  # (n_stddevs, n_bins)
        self.register_buffer('Y', torch.stack(Y, dim=0))  # (n_stddevs, n_bins)

    def sample(self, std_idx):
        """
        Args:
            std_idx:  Indices of standard deviation.
        Returns:
            samples:  Angular samples [0, PI), same size as std.
        """
        size = std_idx.size()
        std_idx = std_idx.flatten() # (N,)
        
        # Samples from histogram
        prob = self.Y[std_idx]  # (N, n_bins)
        bin_idx = torch.multinomial(prob[:, :-1], num_samples=1).squeeze(-1)    # (N,)
        bin_start = self.X[std_idx, bin_idx]    # (N,)
        bin_width = self.X[std_idx, bin_idx+1] - self.X[std_idx, bin_idx]
        samples_hist = bin_start + torch.rand_like(bin_start) * bin_width    # (N,)

        # Samples from Gaussian approximation
        mean_gaussian = self.stddevs[std_idx]*2
        std_gaussian = self.stddevs[std_idx]
        samples_gaussian = mean_gaussian + torch.randn_like(mean_gaussian) * std_gaussian
        samples_gaussian = samples_gaussian.abs() % math.pi

        # Choose from histogram or Gaussian
        gaussian_flag = self.approx_flag[std_idx]
        samples = torch.where(gaussian_flag, samples_gaussian, samples_hist)

        return samples.reshape(size)


def random_normal_so3(std_idx, angular_distrib, device='cpu'):
    size = std_idx.size()
    u = F.normalize(torch.randn(list(size)+[3,], device=device), dim=-1)
    theta = angular_distrib.sample(std_idx)
    w = u * theta[..., None]
    return w

AbDock/src/modules/common/so3.py
pengzhangzhi/ab_opt
import torch
from torch.nn import Module, Linear, LayerNorm, Sequential, ReLU

from ..common.geometry import compose_rotation_and_translation, quaternion_to_rotation_matrix, repr_6d_to_rotation_matrix


class FrameRotationTranslationPrediction(Module):

    def __init__(self, feat_dim, rot_repr, nn_type='mlp'):
        super().__init__()
        assert rot_repr in ('quaternion', '6d')
        self.rot_repr = rot_repr
        if rot_repr == 'quaternion':
            out_dim = 3 + 3
        elif rot_repr == '6d':
            out_dim = 6 + 3
        
        if nn_type == 'linear':
            self.nn = Linear(feat_dim, out_dim)
        elif nn_type == 'mlp':
            self.nn = Sequential(
                Linear(feat_dim, feat_dim), ReLU(),
                Linear(feat_dim, feat_dim), ReLU(),
                Linear(feat_dim, out_dim)
            )
        else:
            raise ValueError('Unknown nn_type: %s' % nn_type)

    def forward(self, x):
        y = self.nn(x)  # (..., d+3)
        if self.rot_repr == 'quaternion':
            quaternion = torch.cat([torch.ones_like(y[..., :1]), y[..., 0:3]], dim=-1)
            R_delta = quaternion_to_rotation_matrix(quaternion)
            t_delta = y[..., 3:6]
            return R_delta, t_delta
        elif self.rot_repr == '6d':
            R_delta = repr_6d_to_rotation_matrix(y[..., 0:6])
            t_delta = y[..., 6:9]
            return R_delta, t_delta


class FrameUpdate(Module):

    def __init__(self, node_feat_dim, rot_repr='quaternion', rot_tran_nn_type='mlp'):
        super().__init__()
        self.transition_mlp = Sequential(
            Linear(node_feat_dim, node_feat_dim), ReLU(),
            Linear(node_feat_dim, node_feat_dim), ReLU(),
            Linear(node_feat_dim, node_feat_dim),
        )
        self.transition_layer_norm = LayerNorm(node_feat_dim)

        self.rot_tran = FrameRotationTranslationPrediction(node_feat_dim, rot_repr, nn_type=rot_tran_nn_type)
    
    def forward(self, R, t, x, mask_generate):
        """
        Args:
            R:  Frame basis matrices, (N, L, 3, 3_index).
            t:  Frame external (absolute) coordinates, (N, L, 3). Unit: Angstrom.
            x:  Node-wise features, (N, L, F).
            mask_generate:   Masks, (N, L).
        Returns:
            R': Updated basis matrices, (N, L, 3, 3_index).
            t': Updated coordinates, (N, L, 3).
        """
        x = self.transition_layer_norm(x + self.transition_mlp(x))

        R_delta, t_delta = self.rot_tran(x) # (N, L, 3, 3), (N, L, 3)
        R_new, t_new = compose_rotation_and_translation(R, t, R_delta, t_delta)

        mask_R = mask_generate[:, :, None, None].expand_as(R)
        mask_t = mask_generate[:, :, None].expand_as(t)

        R_new = torch.where(mask_R, R_new, R)
        t_new = torch.where(mask_t, t_new, t)

        return R_new, t_new

AbDock/src/modules/common/structure.py
pengzhangzhi/ab_opt
import torch
import torch.nn.functional as F


def get_consecutive_flag(chain_nb, res_nb, mask):
    """
    Args:
        chain_nb, res_nb
    Returns:
        consec: A flag tensor indicating whether residue-i is connected to residue-(i+1), 
                BoolTensor, (B, L-1)[b, i].
    """
    d_res_nb = (res_nb[:, 1:] - res_nb[:, :-1]).abs()   # (B, L-1)
    same_chain = (chain_nb[:, 1:] == chain_nb[:, :-1])
    consec = torch.logical_and(d_res_nb == 1, same_chain)
    consec = torch.logical_and(consec, mask[:, :-1])
    return consec


def get_terminus_flag(chain_nb, res_nb, mask):
    consec = get_consecutive_flag(chain_nb, res_nb, mask)
    N_term_flag = F.pad(torch.logical_not(consec), pad=(1, 0), value=1)
    C_term_flag = F.pad(torch.logical_not(consec), pad=(0, 1), value=1)
    return N_term_flag, C_term_flag

AbDock/src/modules/common/topology.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F
import functools
from tqdm.auto import tqdm

from src.modules.common.geometry import apply_rotation_to_vector, quaternion_1ijk_to_rotation_matrix
from src.modules.common.plddt import PerResidueLDDTCaPredictor, compute_plddt, lddt_loss
from src.modules.common.prmsd import PerResidueRMSDCaPredictor, pRMSDCa
from src.modules.common.so3 import so3vec_to_rotation, rotation_to_so3vec, random_uniform_so3
from src.modules.encoders.ga import GAEncoder
from .transition import RotationTransition, PositionTransition, AminoacidCategoricalTransition


def rotation_matrix_cosine_loss(R_pred, R_true):
    """
    Args:
        R_pred: (*, 3, 3).
        R_true: (*, 3, 3).
    Returns:
        Per-matrix losses, (*, ).
    """
    size = list(R_pred.shape[:-2])
    ncol = R_pred.numel() // 3

    RT_pred = R_pred.transpose(-2, -1).reshape(ncol, 3) # (ncol, 3)
    RT_true = R_true.transpose(-2, -1).reshape(ncol, 3) # (ncol, 3)

    ones = torch.ones([ncol, ], dtype=torch.long, device=R_pred.device)
    loss = F.cosine_embedding_loss(RT_pred, RT_true, ones, reduction='none')  # (ncol*3, )
    loss = loss.reshape(size + [3]).sum(dim=-1)    # (*, )
    return loss


class EpsilonNet(nn.Module):

    def __init__(self, res_feat_dim, pair_feat_dim, num_layers, no_bins, encoder_opt={}):
        super().__init__()
        self.current_sequence_embedding = nn.Embedding(25, res_feat_dim)  # 22 is padding
        self.res_feat_mixer = nn.Sequential(
            nn.Linear(res_feat_dim * 2, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim),
        )
        self.encoder = GAEncoder(res_feat_dim, pair_feat_dim, num_layers, **encoder_opt)

        self.eps_crd_net = nn.Sequential(
            nn.Linear(res_feat_dim+3, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, 3)
        )

        self.eps_rot_net = nn.Sequential(
            nn.Linear(res_feat_dim+3, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, 3)
        )

        self.eps_seq_net = nn.Sequential(
            nn.Linear(res_feat_dim+3, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, res_feat_dim), nn.ReLU(),
            nn.Linear(res_feat_dim, 20), nn.Softmax(dim=-1) 
        )
        self.prmsd_predictor = PerResidueRMSDCaPredictor(
            no_bins, res_feat_dim+3, res_feat_dim,
        )
        # self.plddt_predictor = PerResidueLDDTCaPredictor(
        #     no_bins, res_feat_dim, res_feat_dim,
        # )

    def forward(self, v_t, p_t, s_t, res_feat, pair_feat, beta, mask_generate, mask_res):
        """
        Args:
            v_t:    (N, L, 3).
            p_t:    (N, L, 3).
            s_t:    (N, L).
            res_feat:   (N, L, res_dim).
            pair_feat:  (N, L, L, pair_dim).
            beta:   (N,).
            mask_generate:    (N, L).
            mask_res:       (N, L).
        Returns:
            v_next: UPDATED (not epsilon) SO3-vector of orietnations, (N, L, 3).
            eps_pos: (N, L, 3).
        """
        N, L = mask_res.size()
        R = so3vec_to_rotation(v_t) # (N, L, 3, 3)

        # s_t = s_t.clamp(min=0, max=19)  # TODO: clamping is good but ugly.
        res_feat = self.res_feat_mixer(torch.cat([res_feat, self.current_sequence_embedding(s_t)], dim=-1)) # [Important] Incorporate sequence at the current step.
        res_feat = self.encoder(R, p_t, res_feat, pair_feat, mask_res)

        t_embed = torch.stack([beta, torch.sin(beta), torch.cos(beta)], dim=-1)[:, None, :].expand(N, L, 3)
        in_feat = torch.cat([res_feat, t_embed], dim=-1)

        # Position changes
        eps_crd = self.eps_crd_net(in_feat)    # (N, L, 3)
        eps_pos = apply_rotation_to_vector(R, eps_crd)  # (N, L, 3)
        eps_pos = torch.where(mask_generate[:, :, None].expand_as(eps_pos), eps_pos, torch.zeros_like(eps_pos))

        # New orientation
        eps_rot = self.eps_rot_net(in_feat)    # (N, L, 3)
        U = quaternion_1ijk_to_rotation_matrix(eps_rot) # (N, L, 3, 3)
        R_next = R @ U
        v_next = rotation_to_so3vec(R_next)     # (N, L, 3)
        v_next = torch.where(mask_generate[:, :, None].expand_as(v_next), v_next, v_t)

        # New sequence categorical distributions
        c_denoised = self.eps_seq_net(in_feat)  # Already softmax-ed, (N, L, 20)
        prmsd_logits = self.prmsd_predictor(in_feat)  # (N, L, no_bins)
        prmsd_logits = prmsd_logits.mean(dim=1)  # (N, no_bins)
        # prmsd_logits = self.plddt_predictor(res_feat)  # (N, L, no_bins)
        return v_next, R_next, eps_pos, c_denoised, prmsd_logits


class FullDPM(nn.Module):

    def __init__(
        self, 
        res_feat_dim, 
        pair_feat_dim, 
        num_steps, 
        eps_net_opt={}, 
        trans_rot_opt={}, 
        trans_pos_opt={}, 
        trans_seq_opt={},
        position_mean=[0.0, 0.0, 0.0],
        position_scale=[10.0],
        obj = 'pred_noise',
        num_bins=20,
        dist_min=0.5,
        dist_max=19.5
    ):
        super().__init__()
        self.eps_net = EpsilonNet(res_feat_dim, pair_feat_dim, **eps_net_opt, no_bins=num_bins)
        self.num_steps = num_steps
        self.trans_rot = RotationTransition(num_steps, **trans_rot_opt)
        self.trans_pos = PositionTransition(num_steps, **trans_pos_opt)
        self.trans_seq = AminoacidCategoricalTransition(num_steps, **trans_seq_opt)

        self.register_buffer('position_mean', torch.FloatTensor(position_mean).view(1, 1, -1))
        self.register_buffer('position_scale', torch.FloatTensor(position_scale).view(1, 1, -1))
        self.register_buffer('_dummy', torch.empty([0, ]))
        self.obj = obj
        assert self.obj in ['pred_x0', 'pred_noise']
        self.prmsd = pRMSDCa(num_bins, dist_min=dist_min, dist_max=dist_max)
    
    
    def _normalize_position(self, p):
        p_norm = (p - self.position_mean) / self.position_scale
        return p_norm

    def _unnormalize_position(self, p_norm):
        p = p_norm * self.position_scale + self.position_mean
        return p

    def forward(self, v_0, p_0, s_0, res_feat, pair_feat, mask_generate, mask_res, denoise_structure, denoise_sequence, t=None):
        N, L = res_feat.shape[:2]
        if t == None:
            t = torch.randint(0, self.num_steps, (N,), dtype=torch.long, device=self._dummy.device)
        p_0 = self._normalize_position(p_0)

        if denoise_structure:
            # Add noise to rotation
            R_0 = so3vec_to_rotation(v_0)
            v_noisy, _ = self.trans_rot.add_noise(v_0, mask_generate, t)
            # Add noise to positions
            p_noisy, eps_p = self.trans_pos.add_noise(p_0, mask_generate, t)
        else:
            R_0 = so3vec_to_rotation(v_0)
            v_noisy = v_0.clone()
            p_noisy = p_0.clone()
            eps_p = torch.zeros_like(p_noisy)

        if denoise_sequence:
            # Add noise to sequence
            _, s_noisy = self.trans_seq.add_noise(s_0, mask_generate, t)
        else:
            s_noisy = s_0.clone()

        beta = self.trans_pos.var_sched.betas[t]
        v_pred, R_pred, p_pred, c_denoised, prmsd_logits = self.eps_net(
            v_noisy, p_noisy, s_noisy, res_feat, pair_feat, beta, mask_generate, mask_res
        )   # (N, L, 3), (N, L, 3, 3), (N, L, 3), (N, L, 20), (N, L)
        if self.obj == 'pred_x0':
            p_true = p_0
            pred_p0 = p_pred
        elif self.obj == 'pred_noise':
            p_true = p_noisy
            pred_p0 = self.trans_pos.pred_start_from_noise(p_0, p_pred,mask_generate, t)
        loss_dict = {}
        
        # PRMSD loss
        # per_rmsd = pRMSDCa.calc_per_rmsd(self._unnormalize_position(pred_p0), self._unnormalize_position(p_0))
        rmsd = pRMSDCa.calc_rmsd(self._unnormalize_position(pred_p0), self._unnormalize_position(p_0), mask_generate)
        prmsd_loss = self.prmsd(prmsd_logits, rmsd.detach(), mask_generate[:, 0])
        # prmsd_loss = self.prmsd(prmsd_logits, per_rmsd.detach(), mask_generate)
        loss_dict['prmsd'] = prmsd_loss
        
        # # pLDDT loss
        # plddt_loss = lddt_loss(
        #     prmsd_logits, 
        #     self._unnormalize_position(pred_p0), self._unnormalize_position(p_0), 
        #     mask_generate.unsqueeze(-1)
        # )
        # loss_dict['plddt'] = plddt_loss
        
        # distance loss
        if self.obj == 'pred_x0':
            dist_loss = calc_dist_loss(p_pred, p_true, mask_generate, mask_res)
            loss_dict['dist'] = dist_loss
        
        # Rotation loss
        loss_rot = rotation_matrix_cosine_loss(R_pred, R_0) # (N, L)
        loss_rot = (loss_rot * mask_generate).sum() / (mask_generate.sum().float() + 1e-8)
        loss_dict['rot'] = loss_rot

        # Position loss
        loss_pos = F.mse_loss(p_pred, p_true, reduction='none').sum(dim=-1)  # (N, L)
        loss_pos = (loss_pos * mask_generate).sum() / (mask_generate.sum().float() + 1e-8)
        loss_dict['pos'] = loss_pos

        # Sequence categorical loss
        post_true = self.trans_seq.posterior(s_noisy, s_0, t)
        log_post_pred = torch.log(self.trans_seq.posterior(s_noisy, c_denoised, t) + 1e-8)
        kldiv = F.kl_div(
            input=log_post_pred, 
            target=post_true, 
            reduction='none',
            log_target=False
        ).sum(dim=-1)    # (N, L)
        loss_seq = (kldiv * mask_generate).sum() / (mask_generate.sum().float() + 1e-8)
        loss_dict['seq'] = loss_seq

        return loss_dict

    @torch.no_grad()
    def sample(
        self, 
        v, p, s, 
        res_feat, pair_feat, 
        mask_generate, mask_res, 
        sample_structure=True, sample_sequence=True,
        pbar=False, **kwargs
    ):
        """
        Args:
            v:  Orientations of contextual residues, (N, L, 3).
            p:  Positions of contextual residues, (N, L, 3).
            s:  Sequence of contextual residues, (N, L).
        """
        N, L = v.shape[:2]
        p = self._normalize_position(p)

        # Set the orientation and position of residues to be predicted to random values
        if sample_structure:
            v_rand = random_uniform_so3([N, L], device=self._dummy.device)
            p_rand = torch.randn_like(p)
            v_init = torch.where(mask_generate[:, :, None].expand_as(v), v_rand, v)
            p_init = torch.where(mask_generate[:, :, None].expand_as(p), p_rand, p)
        else:
            v_init, p_init = v, p

        if sample_sequence:
            s_rand = torch.randint_like(s, low=0, high=19)
            s_init = torch.where(mask_generate, s_rand, s)
        else:
            s_init = s

        traj = {self.num_steps: [v_init, self._unnormalize_position(p_init), s_init, torch.zeros_like(s_init), torch.ones_like(s_init)]}
        if pbar:
            pbar = functools.partial(tqdm, total=self.num_steps, desc='Sampling')
        else:
            pbar = lambda x: x
        for t in pbar(range(self.num_steps, 0, -1)):
            v_t, p_t, s_t,_,__ = traj[t]
            p_t = self._normalize_position(p_t)
            
            beta = self.trans_pos.var_sched.betas[t].expand([N, ])
            t_tensor = torch.full([N, ], fill_value=t, dtype=torch.long, device=self._dummy.device)

            v_next, R_next, p_pred, c_denoised, prmsd_logits = self.eps_net(
                v_t, p_t, s_t, res_feat, pair_feat, beta, mask_generate, mask_res
            )   # (N, L, 3), (N, L, 3, 3), (N, L, 3)
            prmsd_score = self.prmsd.compute_prmsd(prmsd_logits)
            # plddt_score = compute_plddt(prmsd_logits)
            if self.obj == 'pred_x0':
                eps_p = self.trans_pos.pred_noise_from_start(p_t, p_pred, mask_generate, t_tensor)
            elif self.obj == 'pred_noise':
                eps_p = p_pred 
            v_next = self.trans_rot.denoise(v_t, v_next, mask_generate, t_tensor)
            p_next = self.trans_pos.denoise(p_t, eps_p, mask_generate, t_tensor)
            logits, s_next = self.trans_seq.denoise(s_t, c_denoised, mask_generate, t_tensor)
            perplexity = calc_perplexity(logits,mask_generate)
            if not sample_structure:
                v_next, p_next = v_t, p_t
            if not sample_sequence:
                s_next = s_t

            traj[t-1] = [v_next, self._unnormalize_position(p_next), s_next]+[prmsd_score.cpu(), perplexity.cpu()]
            traj[t] = list(x.cpu() for x in traj[t])    # Move previous states to cpu memory.

        return traj

    @torch.no_grad()
    def optimize(
        self, 
        v, p, s, 
        opt_step: int,
        res_feat, pair_feat, 
        mask_generate, mask_res, 
        sample_structure=True, sample_sequence=True,
        pbar=False,
    ):
        """
        Description:
            First adds noise to the given structure, then denoises it.
        """
        N, L = v.shape[:2]
        p = self._normalize_position(p)
        t = torch.full([N, ], fill_value=opt_step, dtype=torch.long, device=self._dummy.device)

        # Set the orientation and position of residues to be predicted to random values
        if sample_structure:
            # Add noise to rotation
            v_noisy, _ = self.trans_rot.add_noise(v, mask_generate, t)
            # Add noise to positions
            p_noisy, _ = self.trans_pos.add_noise(p, mask_generate, t)
            v_init = torch.where(mask_generate[:, :, None].expand_as(v), v_noisy, v)
            p_init = torch.where(mask_generate[:, :, None].expand_as(p), p_noisy, p)
        else:
            v_init, p_init = v, p

        if sample_sequence:
            _, s_noisy = self.trans_seq.add_noise(s, mask_generate, t)
            s_init = torch.where(mask_generate, s_noisy, s)
        else:
            s_init = s

        traj = {opt_step: (v_init, self._unnormalize_position(p_init), s_init, torch.zeros_like(s_init), torch.ones_like(s_init))}
        if pbar:
            pbar = functools.partial(tqdm, total=opt_step, desc='Optimizing')
        else:
            pbar = lambda x: x
        for t in pbar(range(opt_step, 0, -1)):
            v_t, p_t, s_t,_,__ = traj[t]
            p_t = self._normalize_position(p_t)
            
            beta = self.trans_pos.var_sched.betas[t].expand([N, ])
            t_tensor = torch.full([N, ], fill_value=t, dtype=torch.long, device=self._dummy.device)

            v_next, R_next, eps_p, c_denoised,prmsd_logits = self.eps_net(
                v_t, p_t, s_t, res_feat, pair_feat, beta, mask_generate, mask_res
            )   # (N, L, 3), (N, L, 3, 3), (N, L, 3)
            prmsd_score = self.prmsd.compute_prmsd(prmsd_logits)
            v_next = self.trans_rot.denoise(v_t, v_next, mask_generate, t_tensor)
            p_next = self.trans_pos.denoise(p_t, eps_p, mask_generate, t_tensor)
            logits, s_next = self.trans_seq.denoise(s_t, c_denoised, mask_generate, t_tensor)
            perplexity = calc_perplexity(logits)
            if not sample_structure:
                v_next, p_next = v_t, p_t
            if not sample_sequence:
                s_next = s_t

            traj[t-1] = (v_next, self._unnormalize_position(p_next), s_next, prmsd_score, perplexity)
            traj[t] = tuple(x.cpu() for x in traj[t])    # Move previous states to cpu memory.

        return traj

def calc_dist_loss(p_pred, p_true, mask_generate, mask_res):
    pred_dist_map = torch.cdist(p_pred, p_pred)   # (N, L, L)
    true_dist_map = torch.cdist(p_true, p_true)   # (N, L, L)
    # repeat mask_generate to match the shape of distance map
    mask_map = mask_res[:, :, None] & mask_res[:, None, :]  # (N, L, L)
    mask_generate_expand = mask_generate[:, :, None].expand_as(pred_dist_map) & mask_map # (N, L, L)
    pred_dist = torch.masked_select(pred_dist_map, mask_generate_expand)
    true_dist = torch.masked_select(true_dist_map, mask_generate_expand)
    loss = nn.SmoothL1Loss(reduction='none')(pred_dist, true_dist).mean()
    return loss

def calc_perplexity(logits, mask_generate=None):
    """
    calculate the perplexity of a sequence of logits

    Args:
        logits (torch.Tensor): (*, L, V)
        mask_generate (torch.Tensor): (*, L)

    Returns:
        perplexity (torch.Tensor): (*, )
    """
    if mask_generate is None:
        mask_generate = torch.ones_like(logits[..., 0], dtype=torch.bool)
    max_probs = F.softmax(logits, dim=-1).max(dim=-1)[0]
    # log_probs = F.log_softmax(logits, dim=-1)
    # max_probs = -1 * torch.max(log_probs, dim=-1)[0]
    max_probs = max_probs * mask_generate.float()
    perplexity = max_probs.sum(dim=-1)
    perplexity = perplexity / mask_generate.float().sum(dim=-1)
    return perplexity
AbDock/src/modules/diffusion/dpm_full.py
pengzhangzhi/ab_opt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from src.modules.common.layers import clampped_one_hot
from src.modules.common.so3 import ApproxAngularDistribution, random_normal_so3, so3vec_to_rotation, rotation_to_so3vec


class VarianceSchedule(nn.Module):

    def __init__(self, num_steps=100, s=0.01):
        super().__init__()
        T = num_steps
        t = torch.arange(0, num_steps+1, dtype=torch.float)
        f_t = torch.cos( (np.pi / 2) * ((t/T) + s) / (1 + s) ) ** 2
        alpha_bars = f_t / f_t[0]

        betas = 1 - (alpha_bars[1:] / alpha_bars[:-1])
        betas = torch.cat([torch.zeros([1]), betas], dim=0)
        betas = betas.clamp_max(0.999)

        sigmas = torch.zeros_like(betas)
        for i in range(1, betas.size(0)):
            sigmas[i] = ((1 - alpha_bars[i-1]) / (1 - alpha_bars[i])) * betas[i]
        sigmas = torch.sqrt(sigmas)
        
        self.register_buffer('betas', betas)
        self.register_buffer('alpha_bars', alpha_bars)
        self.register_buffer('alphas', 1 - betas)
        self.register_buffer('sigmas', sigmas)
        # calculate X0 = sqrt_recip_alphas_cumprod * Xt - sqrt_recipm1_alphas_cumprod * noise
        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alpha_bars))
        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alpha_bars - 1))

class PositionTransition(nn.Module):

    def __init__(self, num_steps, var_sched_opt={}):
        super().__init__()
        self.var_sched = VarianceSchedule(num_steps, **var_sched_opt)
        
    def pred_noise_from_start(self, p_t, p_0, mask_generate, t):
        
        nb_non_batch_dims = len(p_t.shape) - 1
        sqrt_recip_alphas_cumprod = self.var_sched.sqrt_recip_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        sqrt_recipm1_alphas_cumprod = self.var_sched.sqrt_recipm1_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        eps_p = (sqrt_recip_alphas_cumprod * p_t - p_0) / sqrt_recipm1_alphas_cumprod

        eps_p = torch.where(mask_generate[..., None].expand_as(p_t), eps_p, p_t)
        return eps_p
    
    def pred_start_from_noise(self, p_t, eps_p, mask_generate, t):
        nb_non_batch_dims = len(p_t.shape) - 1
        sqrt_recip_alphas_cumprod = self.var_sched.sqrt_recip_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        sqrt_recipm1_alphas_cumprod = self.var_sched.sqrt_recipm1_alphas_cumprod[t].view(-1, *(nb_non_batch_dims*(1,)))
        
        p_0 = sqrt_recip_alphas_cumprod * p_t - sqrt_recipm1_alphas_cumprod * eps_p

        p_0 = torch.where(mask_generate[..., None].expand_as(p_t), p_0, p_t)
        return p_0
    
    def add_noise(self, p_0, mask_generate, t):
        """
        Args:
            p_0:    (N, L, 3).
            mask_generate:    (N, L).
            t:  (N,).
        """
        alpha_bar = self.var_sched.alpha_bars[t]

        c0 = torch.sqrt(alpha_bar).view(-1, 1, 1)
        c1 = torch.sqrt(1 - alpha_bar).view(-1, 1, 1)

        e_rand = torch.randn_like(p_0)
        p_noisy = c0*p_0 + c1*e_rand
        p_noisy = torch.where(mask_generate[..., None].expand_as(p_0), p_noisy, p_0)

        return p_noisy, e_rand

    def denoise(self, p_t, eps_p, mask_generate, t):
        # IMPORTANT:
        #   clampping alpha is to fix the instability issue at the first step (t=T)
        #   it seems like a problem with the ``improved ddpm''.
        alpha = self.var_sched.alphas[t].clamp_min(
            self.var_sched.alphas[-2]
        )
        alpha_bar = self.var_sched.alpha_bars[t]
        sigma = self.var_sched.sigmas[t].view(-1, 1, 1)

        c0 = ( 1.0 / torch.sqrt(alpha + 1e-8) ).view(-1, 1, 1)
        c1 = ( (1 - alpha) / torch.sqrt(1 - alpha_bar + 1e-8) ).view(-1, 1, 1)

        z = torch.where(
            (t > 1)[:, None, None].expand_as(p_t),
            torch.randn_like(p_t),
            torch.zeros_like(p_t),
        )

        p_next = c0 * (p_t - c1 * eps_p) + sigma * z
        p_next = torch.where(mask_generate[..., None].expand_as(p_t), p_next, p_t)
        return p_next


class RotationTransition(nn.Module):

    def __init__(self, num_steps, var_sched_opt={}, angular_distrib_fwd_opt={}, angular_distrib_inv_opt={}):
        super().__init__()
        self.var_sched = VarianceSchedule(num_steps, **var_sched_opt)

        # Forward (perturb)
        c1 = torch.sqrt(1 - self.var_sched.alpha_bars) # (T,).
        self.angular_distrib_fwd = ApproxAngularDistribution(c1.tolist(), **angular_distrib_fwd_opt)

        # Inverse (generate)
        sigma = self.var_sched.sigmas
        self.angular_distrib_inv = ApproxAngularDistribution(sigma.tolist(), **angular_distrib_inv_opt)

        self.register_buffer('_dummy', torch.empty([0, ]))

    def add_noise(self, v_0, mask_generate, t):
        """
        Args:
            v_0:    (N, L, 3).
            mask_generate:    (N, L).
            t:  (N,).
        """
        N, L = mask_generate.size()
        alpha_bar = self.var_sched.alpha_bars[t]
        c0 = torch.sqrt(alpha_bar).view(-1, 1, 1)
        c1 = torch.sqrt(1 - alpha_bar).view(-1, 1, 1)

        # Noise rotation
        e_scaled = random_normal_so3(t[:, None].expand(N, L), self.angular_distrib_fwd, device=self._dummy.device)    # (N, L, 3)
        e_normal = e_scaled / (c1 + 1e-8)
        E_scaled = so3vec_to_rotation(e_scaled)   # (N, L, 3, 3)

        # Scaled true rotation
        R0_scaled = so3vec_to_rotation(c0 * v_0)  # (N, L, 3, 3)

        R_noisy = E_scaled @ R0_scaled
        v_noisy = rotation_to_so3vec(R_noisy)
        v_noisy = torch.where(mask_generate[..., None].expand_as(v_0), v_noisy, v_0)

        return v_noisy, e_scaled

    def denoise(self, v_t, v_next, mask_generate, t):
        N, L = mask_generate.size()
        e = random_normal_so3(t[:, None].expand(N, L), self.angular_distrib_inv, device=self._dummy.device) # (N, L, 3)
        e = torch.where(
            (t > 1)[:, None, None].expand(N, L, 3),
            e, 
            torch.zeros_like(e) # Simply denoise and don't add noise at the last step
        )
        E = so3vec_to_rotation(e)

        R_next = E @ so3vec_to_rotation(v_next)
        v_next = rotation_to_so3vec(R_next)
        v_next = torch.where(mask_generate[..., None].expand_as(v_next), v_next, v_t)

        return v_next


class AminoacidCategoricalTransition(nn.Module):
    
    def __init__(self, num_steps, num_classes=20, var_sched_opt={}):
        super().__init__()
        self.num_classes = num_classes
        self.var_sched = VarianceSchedule(num_steps, **var_sched_opt)

    @staticmethod
    def _sample(c):
        """
        Args:
            c:    (N, L, K).
        Returns:
            x:    (N, L).
        """
        N, L, K = c.size()
        c = c.view(N*L, K) + 1e-8
        x = torch.multinomial(c, 1).view(N, L)
        return x

    def add_noise(self, x_0, mask_generate, t):
        """
        Args:
            x_0:    (N, L)
            mask_generate:    (N, L).
            t:  (N,).
        Returns:
            c_t:    Probability, (N, L, K).
            x_t:    Sample, LongTensor, (N, L).
        """
        N, L = x_0.size()
        K = self.num_classes
        c_0 = clampped_one_hot(x_0, num_classes=K).float() # (N, L, K).
        alpha_bar = self.var_sched.alpha_bars[t][:, None, None] # (N, 1, 1)
        c_noisy = (alpha_bar*c_0) + ( (1-alpha_bar)/K )
        c_t = torch.where(mask_generate[..., None].expand(N,L,K), c_noisy, c_0)
        x_t = self._sample(c_t)
        return c_t, x_t

    def posterior(self, x_t, x_0, t):
        """
        Args:
            x_t:    Category LongTensor (N, L) or Probability FloatTensor (N, L, K).
            x_0:    Category LongTensor (N, L) or Probability FloatTensor (N, L, K).
            t:  (N,).
        Returns:
            theta:  Posterior probability at (t-1)-th step, (N, L, K).
        """
        K = self.num_classes

        if x_t.dim() == 3:
            c_t = x_t   # When x_t is probability distribution.
        else:
            c_t = clampped_one_hot(x_t, num_classes=K).float() # (N, L, K)

        if x_0.dim() == 3:
            c_0 = x_0   # When x_0 is probability distribution.
        else:
            c_0 = clampped_one_hot(x_0, num_classes=K).float() # (N, L, K)

        alpha = self.var_sched.alpha_bars[t][:, None, None]     # (N, 1, 1)
        alpha_bar = self.var_sched.alpha_bars[t][:, None, None] # (N, 1, 1)

        theta = ((alpha*c_t) + (1-alpha)/K) * ((alpha_bar*c_0) + (1-alpha_bar)/K)   # (N, L, K)
        theta = theta / (theta.sum(dim=-1, keepdim=True) + 1e-8)
        return theta

    def denoise(self, x_t, c_0_pred, mask_generate, t):
        """
        Args:
            x_t:        (N, L).
            c_0_pred:   Normalized probability predicted by networks, (N, L, K).
            mask_generate:    (N, L).
            t:  (N,).
        Returns:
            post:   Posterior probability at (t-1)-th step, (N, L, K).
            x_next: Sample at (t-1)-th step, LongTensor, (N, L).
        """
        c_t = clampped_one_hot(x_t, num_classes=self.num_classes).float()  # (N, L, K)
        post = self.posterior(c_t, c_0_pred, t=t)   # (N, L, K)
        post = torch.where(mask_generate[..., None].expand(post.size()), post, c_t)
        x_next = self._sample(post)
        return post, x_next

AbDock/src/modules/diffusion/transition.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from src.modules.common.geometry import global_to_local, local_to_global, normalize_vector, construct_3d_basis, angstrom_to_nm
from src.modules.common.layers import mask_zero, LayerNorm
from src.utils.protein.constants import BBHeavyAtom


def _alpha_from_logits(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(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)


class GABlock(nn.Module):

    def __init__(self, node_feat_dim, pair_feat_dim, value_dim=32, query_key_dim=32, num_query_points=8,
                 num_value_points=8, num_heads=12, bias=False):
        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=bias)
        self.proj_key = nn.Linear(node_feat_dim, query_key_dim * num_heads, bias=bias)
        self.proj_value = nn.Linear(node_feat_dim, value_dim * num_heads, bias=bias)

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

        # Spatial
        self.spatial_coef = nn.Parameter(torch.full([1, 1, 1, self.num_heads], fill_value=np.log(np.exp(1.) - 1.)),
                                         requires_grad=True)
        self.proj_query_point = nn.Linear(node_feat_dim, num_query_points * num_heads * 3, bias=bias)
        self.proj_key_point = nn.Linear(node_feat_dim, num_query_points * num_heads * 3, bias=bias)
        self.proj_value_point = nn.Linear(node_feat_dim, num_value_points * num_heads * 3, bias=bias)

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

        self.layer_norm_1 = LayerNorm(node_feat_dim)
        self.mlp_transition = nn.Sequential(nn.Linear(node_feat_dim, node_feat_dim), nn.ReLU(),
                                            nn.Linear(node_feat_dim, node_feat_dim), nn.ReLU(),
                                            nn.Linear(node_feat_dim, node_feat_dim))
        self.layer_norm_2 = LayerNorm(node_feat_dim)

    def _node_logits(self, x):
        query_l = _heads(self.proj_query(x), self.num_heads, self.query_key_dim)  # (N, L, n_heads, qk_ch)
        key_l = _heads(self.proj_key(x), self.num_heads, self.query_key_dim)  # (N, L, n_heads, qk_ch)
        logits_node = (query_l.unsqueeze(2) * key_l.unsqueeze(1) *
                       (1 / np.sqrt(self.query_key_dim))).sum(-1)  # (N, L, L, num_heads)
        return logits_node

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

    def _spatial_logits(self, R, t, x):
        N, L, _ = t.size()

        # Query
        query_points = _heads(self.proj_query_point(x), self.num_heads * self.num_query_points,
                              3)  # (N, L, n_heads * n_pnts, 3)
        query_points = local_to_global(R, t, query_points)  # Global query coordinates, (N, L, n_heads * n_pnts, 3)
        query_s = query_points.reshape(N, L, self.num_heads, -1)  # (N, L, n_heads, n_pnts*3)

        # Key
        key_points = _heads(self.proj_key_point(x), self.num_heads * self.num_query_points,
                            3)  # (N, L, 3, n_heads * n_pnts)
        key_points = local_to_global(R, t, key_points)  # Global key coordinates, (N, L, n_heads * n_pnts, 3)
        key_s = key_points.reshape(N, L, self.num_heads, -1)  # (N, L, n_heads, n_pnts*3)

        # Q-K Product
        sum_sq_dist = ((query_s.unsqueeze(2) - key_s.unsqueeze(1)) ** 2).sum(-1)  # (N, L, L, n_heads)
        gamma = F.softplus(self.spatial_coef)
        logits_spatial = sum_sq_dist * ((-1 * gamma * np.sqrt(2 / (9 * self.num_query_points)))
                                        / 2)  # (N, L, L, n_heads)
        return logits_spatial

    def _pair_aggregation(self, alpha, z):
        N, L = z.shape[:2]
        feat_p2n = alpha.unsqueeze(-1) * z.unsqueeze(-2)  # (N, L, L, n_heads, C)
        feat_p2n = feat_p2n.sum(dim=2)  # (N, L, n_heads, C)
        return feat_p2n.reshape(N, L, -1)

    def _node_aggregation(self, alpha, x):
        N, L = x.shape[:2]
        value_l = _heads(self.proj_value(x), self.num_heads, self.query_key_dim)  # (N, L, n_heads, v_ch)
        feat_node = alpha.unsqueeze(-1) * value_l.unsqueeze(1)  # (N, L, L, n_heads, *) @ (N, *, L, n_heads, v_ch)
        feat_node = feat_node.sum(dim=2)  # (N, L, n_heads, v_ch)
        return feat_node.reshape(N, L, -1)

    def _spatial_aggregation(self, alpha, R, t, x):
        N, L, _ = t.size()
        value_points = _heads(self.proj_value_point(x), self.num_heads * self.num_value_points,
                              3)  # (N, L, n_heads * n_v_pnts, 3)
        value_points = local_to_global(R, t, value_points.reshape(N, L, self.num_heads, self.num_value_points,
                                                                  3))  # (N, L, n_heads, n_v_pnts, 3)
        aggr_points = alpha.reshape(N, L, L, self.num_heads, 1, 1) * \
                      value_points.unsqueeze(1)  # (N, *, L, n_heads, n_pnts, 3)
        aggr_points = aggr_points.sum(dim=2)  # (N, L, n_heads, n_pnts, 3)

        feat_points = global_to_local(R, t, aggr_points)  # (N, L, n_heads, n_pnts, 3)
        feat_distance = feat_points.norm(dim=-1)  # (N, L, n_heads, n_pnts)
        feat_direction = normalize_vector(feat_points, dim=-1, eps=1e-4)  # (N, L, n_heads, n_pnts, 3)

        feat_spatial = torch.cat([
            feat_points.reshape(N, L, -1),
            feat_distance.reshape(N, L, -1),
            feat_direction.reshape(N, L, -1),
        ], dim=-1)

        return feat_spatial

    def forward(self, R, t, x, z, mask):
        """
        Args:
            R:  Frame basis matrices, (N, L, 3, 3_index).
            t:  Frame external (absolute) coordinates, (N, L, 3).
            x:  Node-wise features, (N, L, F).
            z:  Pair-wise features, (N, L, L, C).
            mask:   Masks, (N, L).
        Returns:
            x': Updated node-wise features, (N, L, F).
        """
        # Attention logits
        logits_node = self._node_logits(x)
        logits_pair = self._pair_logits(z)
        logits_spatial = self._spatial_logits(R, t, x)
        # Summing logits up and apply `softmax`.
        logits_sum = logits_node + logits_pair + logits_spatial
        alpha = _alpha_from_logits(logits_sum * np.sqrt(1 / 3), mask)  # (N, L, L, n_heads)

        # Aggregate features
        feat_p2n = self._pair_aggregation(alpha, z)
        feat_node = self._node_aggregation(alpha, x)
        feat_spatial = self._spatial_aggregation(alpha, R, t, x)

        # Finally
        feat_all = self.out_transform(torch.cat([feat_p2n, feat_node, feat_spatial], dim=-1))  # (N, L, F)
        feat_all = mask_zero(mask.unsqueeze(-1), feat_all)
        x_updated = self.layer_norm_1(x + feat_all)
        x_updated = self.layer_norm_2(x_updated + self.mlp_transition(x_updated))
        return x_updated


class GAEncoder(nn.Module):

    def __init__(self, node_feat_dim, pair_feat_dim, num_layers, ga_block_opt={}):
        super(GAEncoder, self).__init__()
        self.blocks = nn.ModuleList([
            GABlock(node_feat_dim, pair_feat_dim, **ga_block_opt) 
            for _ in range(num_layers)
        ])

    def forward(self, R, t, res_feat, pair_feat, mask):
        for i, block in enumerate(self.blocks):
            res_feat = block(R, t, res_feat, pair_feat, mask)
        return res_feat

AbDock/src/modules/encoders/ga.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn
import torch.nn.functional as F

from src.modules.common.geometry import angstrom_to_nm, pairwise_dihedrals
from src.modules.common.layers import AngularEncoding
from src.utils.protein.constants import BBHeavyAtom, AA


class PairEmbedding(nn.Module):

    def __init__(self, feat_dim, max_num_atoms, max_aa_types=22, max_relpos=32):
        super().__init__()
        self.max_num_atoms = max_num_atoms
        self.max_aa_types = max_aa_types
        self.max_relpos = max_relpos
        self.aa_pair_embed = nn.Embedding(self.max_aa_types*self.max_aa_types, feat_dim)
        self.relpos_embed = nn.Embedding(2*max_relpos+1, feat_dim)

        self.aapair_to_distcoef = nn.Embedding(self.max_aa_types*self.max_aa_types, max_num_atoms*max_num_atoms)
        nn.init.zeros_(self.aapair_to_distcoef.weight)
        self.distance_embed = nn.Sequential(
            nn.Linear(max_num_atoms*max_num_atoms, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim), nn.ReLU(),
        )

        self.dihedral_embed = AngularEncoding()
        feat_dihed_dim = self.dihedral_embed.get_out_dim(2) # Phi and Psi

        infeat_dim = feat_dim+feat_dim+feat_dim+feat_dihed_dim
        self.out_mlp = nn.Sequential(
            nn.Linear(infeat_dim, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim),
        )

    def forward(self, aa, res_nb, chain_nb, pos_atoms, mask_atoms, structure_mask=None, sequence_mask=None):
        """
        Args:
            aa: (N, L).
            res_nb: (N, L).
            chain_nb: (N, L).
            pos_atoms:  (N, L, A, 3)
            mask_atoms: (N, L, A)
            structure_mask: (N, L)
            sequence_mask:  (N, L), mask out unknown amino acids to generate.

        Returns:
            (N, L, L, feat_dim)
        """
        N, L = aa.size()

        # Remove other atoms
        pos_atoms = pos_atoms[:, :, :self.max_num_atoms]
        mask_atoms = mask_atoms[:, :, :self.max_num_atoms]

        mask_residue = mask_atoms[:, :, BBHeavyAtom.CA] # (N, L)
        mask_pair = mask_residue[:, :, None] * mask_residue[:, None, :]
        pair_structure_mask = structure_mask[:, :, None] * structure_mask[:, None, :] if structure_mask is not None else None

        # Pair identities
        if sequence_mask is not None:
            # Avoid data leakage at training time
            aa = torch.where(sequence_mask, aa, torch.full_like(aa, fill_value=AA.UNK))
        aa_pair = aa[:,:,None]*self.max_aa_types + aa[:,None,:]    # (N, L, L)
        feat_aapair = self.aa_pair_embed(aa_pair)
    
        # Relative sequential positions
        same_chain = (chain_nb[:, :, None] == chain_nb[:, None, :])
        relpos = torch.clamp(
            res_nb[:,:,None] - res_nb[:,None,:], 
            min=-self.max_relpos, max=self.max_relpos,
        )   # (N, L, L)
        feat_relpos = self.relpos_embed(relpos + self.max_relpos) * same_chain[:,:,:,None]

        # Distances
        d = angstrom_to_nm(torch.linalg.norm(
            pos_atoms[:,:,None,:,None] - pos_atoms[:,None,:,None,:],
            dim = -1, ord = 2,
        )).reshape(N, L, L, -1) # (N, L, L, A*A)
        c = F.softplus(self.aapair_to_distcoef(aa_pair))    # (N, L, L, A*A)
        d_gauss = torch.exp(-1 * c * d**2)
        mask_atom_pair = (mask_atoms[:,:,None,:,None] * mask_atoms[:,None,:,None,:]).reshape(N, L, L, -1)
        feat_dist = self.distance_embed(d_gauss * mask_atom_pair)
        if pair_structure_mask is not None:
            # Avoid data leakage at training time
            feat_dist = feat_dist * pair_structure_mask[:, :, :, None]

        # Orientations
        dihed = pairwise_dihedrals(pos_atoms)   # (N, L, L, 2)
        feat_dihed = self.dihedral_embed(dihed)
        if pair_structure_mask is not None:
            # Avoid data leakage at training time
            feat_dihed = feat_dihed * pair_structure_mask[:, :, :, None]

        # All
        feat_all = torch.cat([feat_aapair, feat_relpos, feat_dist, feat_dihed], dim=-1)
        feat_all = self.out_mlp(feat_all)   # (N, L, L, F)
        feat_all = feat_all * mask_pair[:, :, :, None]

        return feat_all


AbDock/src/modules/encoders/pair.py
pengzhangzhi/ab_opt
import torch
import torch.nn as nn

from src.modules.common.geometry import construct_3d_basis, global_to_local, get_backbone_dihedral_angles
from src.modules.common.layers import AngularEncoding
from src.utils.protein.constants import BBHeavyAtom, AA


class ResidueEmbedding(nn.Module):

    def __init__(self, feat_dim, max_num_atoms, max_aa_types=22):
        super().__init__()
        self.max_num_atoms = max_num_atoms
        self.max_aa_types = max_aa_types
        self.aatype_embed = nn.Embedding(self.max_aa_types, feat_dim)
        self.dihed_embed = AngularEncoding()
        self.type_embed = nn.Embedding(10, feat_dim, padding_idx=0)    # 1: Heavy, 2: Light, 3: Ag
        infeat_dim = feat_dim + (self.max_aa_types*max_num_atoms*3) + self.dihed_embed.get_out_dim(3) + feat_dim
        self.mlp = nn.Sequential(
            nn.Linear(infeat_dim, feat_dim * 2), nn.ReLU(),
            nn.Linear(feat_dim * 2, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim), nn.ReLU(),
            nn.Linear(feat_dim, feat_dim)
        )

    def forward(self, aa, res_nb, chain_nb, pos_atoms, mask_atoms, fragment_type, structure_mask=None, sequence_mask=None):
        """
        Args:
            aa:         (N, L).
            res_nb:     (N, L).
            chain_nb:   (N, L).
            pos_atoms:  (N, L, A, 3).
            mask_atoms: (N, L, A).
            fragment_type:  (N, L).
            structure_mask: (N, L), mask out unknown structures to generate.
            sequence_mask:  (N, L), mask out unknown amino acids to generate.
        """
        N, L = aa.size()
        mask_residue = mask_atoms[:, :, BBHeavyAtom.CA] # (N, L)

        # Remove other atoms
        pos_atoms = pos_atoms[:, :, :self.max_num_atoms]
        mask_atoms = mask_atoms[:, :, :self.max_num_atoms]

        # Amino acid identity features
        if sequence_mask is not None:
            # Avoid data leakage at training time
            aa = torch.where(sequence_mask, aa, torch.full_like(aa, fill_value=AA.UNK))
        aa_feat = self.aatype_embed(aa) # (N, L, feat)

        # Coordinate features
        R = construct_3d_basis(
            pos_atoms[:, :, BBHeavyAtom.CA], 
            pos_atoms[:, :, BBHeavyAtom.C], 
            pos_atoms[:, :, BBHeavyAtom.N]
        )
        t = pos_atoms[:, :, BBHeavyAtom.CA]
        crd = global_to_local(R, t, pos_atoms)    # (N, L, A, 3)
        crd_mask = mask_atoms[:, :, :, None].expand_as(crd)
        crd = torch.where(crd_mask, crd, torch.zeros_like(crd))

        aa_expand  = aa[:, :, None, None, None].expand(N, L, self.max_aa_types, self.max_num_atoms, 3)
        rng_expand = torch.arange(0, self.max_aa_types)[None, None, :, None, None].expand(N, L, self.max_aa_types, self.max_num_atoms, 3).to(aa_expand)
        place_mask = (aa_expand == rng_expand)
        crd_expand = crd[:, :, None, :, :].expand(N, L, self.max_aa_types, self.max_num_atoms, 3)
        crd_expand = torch.where(place_mask, crd_expand, torch.zeros_like(crd_expand))
        crd_feat = crd_expand.reshape(N, L, self.max_aa_types*self.max_num_atoms*3)
        if structure_mask is not None:
            # Avoid data leakage at training time
            crd_feat = crd_feat * structure_mask[:, :, None]

        # Backbone dihedral features
        bb_dihedral, mask_bb_dihed = get_backbone_dihedral_angles(pos_atoms, chain_nb=chain_nb, res_nb=res_nb, mask=mask_residue)
        dihed_feat = self.dihed_embed(bb_dihedral[:, :, :, None]) * mask_bb_dihed[:, :, :, None]  # (N, L, 3, dihed/3)
        dihed_feat = dihed_feat.reshape(N, L, -1)
        if structure_mask is not None:
            # Avoid data leakage at training time
            dihed_mask = torch.logical_and(
                structure_mask,
                torch.logical_and(
                    torch.roll(structure_mask, shifts=+1, dims=1), 
                    torch.roll(structure_mask, shifts=-1, dims=1)
                ),
            )   # Avoid slight data leakage via dihedral angles of anchor residues
            dihed_feat = dihed_feat * dihed_mask[:, :, None]

        # Type feature
        type_feat = self.type_embed(fragment_type) # (N, L, feat)

        out_feat = self.mlp(torch.cat([aa_feat, crd_feat, dihed_feat, type_feat], dim=-1)) # (N, L, F)
        out_feat = out_feat * mask_residue[:, :, None]
        return out_feat

AbDock/src/modules/encoders/residue.py
pengzhangzhi/ab_opt
import abc
from typing import List


FilePath = str


class DockingEngine(abc.ABC):

    @abc.abstractmethod
    def __enter__(self):
        pass

    @abc.abstractmethod
    def __exit__(self, typ, value, traceback):
        pass

    @abc.abstractmethod
    def set_receptor(self, pdb_path: FilePath):
        pass

    @abc.abstractmethod
    def set_ligand(self, pdb_path: FilePath):
        pass

    @abc.abstractmethod
    def dock(self) -> List[FilePath]:
        pass

AbDock/src/tools/dock/base.py
pengzhangzhi/ab_opt
import os
import shutil
import tempfile
import subprocess
import dataclasses as dc
from typing import List, Optional
from Bio import PDB
from Bio.PDB import Model as PDBModel

from src.tools.renumber import renumber as renumber_chothia
from .base import DockingEngine


def fix_docked_pdb(pdb_path):
    fixed = []
    with open(pdb_path, 'r') as f:
        for ln in f.readlines():
            if (ln.startswith('ATOM') or ln.startswith('HETATM')) and len(ln) == 56:
                fixed.append( ln[:-1] + ' 1.00  0.00              \n' )
            else:
                fixed.append(ln)
    with open(pdb_path, 'w') as f:
        f.write(''.join(fixed))


class HDock(DockingEngine):

    def __init__(
        self, 
        hdock_bin='./bin/hdock',
        createpl_bin='./bin/createpl',
    ):
        super().__init__()
        self.hdock_bin = os.path.realpath(hdock_bin)
        self.createpl_bin = os.path.realpath(createpl_bin)
        self.tmpdir = tempfile.TemporaryDirectory()

        self._has_receptor = False
        self._has_ligand = False

        self._receptor_chains = []
        self._ligand_chains = []

    def __enter__(self):
        return self

    def __exit__(self, typ, value, traceback):
        self.tmpdir.cleanup()

    def set_receptor(self, pdb_path):
        shutil.copyfile(pdb_path, os.path.join(self.tmpdir.name, 'receptor.pdb'))
        self._has_receptor = True

    def set_ligand(self, pdb_path):
        shutil.copyfile(pdb_path, os.path.join(self.tmpdir.name, 'ligand.pdb'))
        self._has_ligand = True

    def _dump_complex_pdb(self):
        parser = PDB.PDBParser(QUIET=True)
        model_receptor = parser.get_structure(None, os.path.join(self.tmpdir.name, 'receptor.pdb'))[0]
        docked_pdb_path = os.path.join(self.tmpdir.name, 'ligand_docked.pdb')
        fix_docked_pdb(docked_pdb_path)
        structure_ligdocked = parser.get_structure(None, docked_pdb_path)

        pdb_io = PDB.PDBIO()
        paths = []
        for i, model_ligdocked in enumerate(structure_ligdocked):
            model_complex = PDBModel.Model(0)
            for chain in model_receptor:
                model_complex.add(chain.copy())
            for chain in model_ligdocked:
                model_complex.add(chain.copy())
            pdb_io.set_structure(model_complex)
            save_path = os.path.join(self.tmpdir.name, f"complex_{i}.pdb")
            pdb_io.save(save_path)
            paths.append(save_path)
        return paths

    def dock(self):
        if not (self._has_receptor and self._has_ligand):
            raise ValueError('Missing receptor or ligand.')
        subprocess.run(
            [self.hdock_bin, "receptor.pdb", "ligand.pdb"],
            cwd=self.tmpdir.name, check=True
        )
        subprocess.run(
            [self.createpl_bin, "Hdock.out", "ligand_docked.pdb"], 
            cwd=self.tmpdir.name, check=True
        )
        return self._dump_complex_pdb()


@dc.dataclass
class DockSite:
    chain: str
    resseq: int


class HDockAntibody(HDock):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._heavy_chain_id = None
        self._epitope_sites: Optional[List[DockSite]] = None

    def set_ligand(self, pdb_path):
        raise NotImplementedError('Please use set_antibody')
    
    def set_receptor(self, pdb_path):
        raise NotImplementedError('Please use set_antigen')

    def set_antigen(self, pdb_path, epitope_sites: Optional[List[DockSite]]=None):
        super().set_receptor(pdb_path)
        self._epitope_sites = epitope_sites

    def set_antibody(self, pdb_path):
        heavy_chains, _ = renumber_chothia(pdb_path, os.path.join(self.tmpdir.name, 'ligand.pdb'))
        self._has_ligand = True
        self._heavy_chain_id = heavy_chains[0]

    def _prepare_lsite(self):
        lsite_content = f"95-102:{self._heavy_chain_id}\n"  # Chothia CDR H3
        with open(os.path.join(self.tmpdir.name, 'lsite.txt'), 'w') as f:
            f.write(lsite_content)
        print(f"[INFO] lsite content: {lsite_content}")

    def _prepare_rsite(self):
        rsite_content = ""
        for site in self._epitope_sites:
            rsite_content += f"{site.resseq}:{site.chain}\n"
        with open(os.path.join(self.tmpdir.name, 'rsite.txt'), 'w') as f:
            f.write(rsite_content)
        print(f"[INFO] rsite content: {rsite_content}")

    def dock(self):
        if not (self._has_receptor and self._has_ligand):
            raise ValueError('Missing receptor or ligand.')
        self._prepare_lsite()

        cmd_hdock = [self.hdock_bin, "receptor.pdb", "ligand.pdb", "-lsite", "lsite.txt"]
        if self._epitope_sites is not None:
            self._prepare_rsite()
            cmd_hdock += ["-rsite", "rsite.txt"]
        subprocess.run(
            cmd_hdock,
            cwd=self.tmpdir.name, check=True
        )

        cmd_pl = [self.createpl_bin, "Hdock.out", "ligand_docked.pdb", "-lsite", "lsite.txt"]
        if self._epitope_sites is not None:
            self._prepare_rsite()
            cmd_pl += ["-rsite", "rsite.txt"]
        subprocess.run(
            cmd_pl, 
            cwd=self.tmpdir.name, check=True
        )
        return self._dump_complex_pdb()


if __name__ == '__main__':
    with HDockAntibody('hdock', 'createpl') as dock:
        dock.set_antigen('./data/dock/receptor.pdb', [DockSite('A', 991)])
        dock.set_antibody('./data/example_dock/3qhf_fv.pdb')
        print(dock.dock())

AbDock/src/tools/dock/hdock.py
pengzhangzhi/ab_opt
import os
import re
import json
import shelve
from Bio import PDB
from typing import Optional, Tuple, List
from dataclasses import dataclass, field


@dataclass
class EvalTask:
    in_path: str
    ref_path: str
    info: dict
    structure: str
    name: str
    method: str
    cdr: str
    ab_chains: List

    residue_first: Optional[Tuple] = None
    residue_last: Optional[Tuple] = None
    
    scores: dict = field(default_factory=dict)

    def get_gen_biopython_model(self):
        parser = PDB.PDBParser(QUIET=True)
        return parser.get_structure(self.in_path, self.in_path)[0]

    def get_ref_biopython_model(self):
        parser = PDB.PDBParser(QUIET=True)
        return parser.get_structure(self.ref_path, self.ref_path)[0]

    def save_to_db(self, db: shelve.Shelf):
        db[self.in_path] = self

    def to_report_dict(self):
        return {
            'method': self.method,
            'structure': self.structure,
            'cdr': self.cdr,
            'filename': os.path.basename(self.in_path),
            **self.scores
        }


class TaskScanner:

    def __init__(self, root, postfix=None, db: Optional[shelve.Shelf]=None):
        super().__init__()
        self.root = root
        self.postfix = postfix
        self.visited = set()
        self.db = db
        if db is not None:
            for k in db.keys():
                self.visited.add(k)

    def _get_metadata(self, fpath):
        json_path = os.path.join(
            os.path.dirname(os.path.dirname(fpath)), 
            'metadata.json'
        )
        tag_name = os.path.basename(os.path.dirname(fpath))
        method_name = os.path.basename(
            os.path.dirname(os.path.dirname(os.path.dirname(fpath)))
        )
        try:
            antibody_chains = set()
            info = None
            with open(json_path, 'r') as f:
                metadata = json.load(f)
            for item in metadata['items']:
                if item['tag'] == tag_name:
                    info = item
                antibody_chains.add(item['residue_first'][0])
            if info is not None:
                info['antibody_chains'] = list(antibody_chains)
                info['structure'] = metadata['identifier']
                info['method'] = method_name
            return info
        except (json.JSONDecodeError, FileNotFoundError) as e:
            return None

    def scan(self) -> List[EvalTask]: 
        tasks = []
        if self.postfix is None or not self.postfix:
            input_fname_pattern = '^\d+\.pdb$'
            ref_fname = 'REF1.pdb'
        else:
            input_fname_pattern = f'^\d+\_{self.postfix}\.pdb$'
            ref_fname = f'REF1_{self.postfix}.pdb'
        for parent, _, files in os.walk(self.root):
            for fname in files:
                fpath = os.path.join(parent, fname)
                if not re.match(input_fname_pattern, fname):
                    continue
                if os.path.getsize(fpath) == 0:
                    continue
                if fpath in self.visited:
                    continue

                # Path to the reference structure
                ref_path = os.path.join(parent, ref_fname)
                if not os.path.exists(ref_path):
                    continue

                # CDR information
                info = self._get_metadata(fpath)
                if info is None:
                    continue
                tasks.append(EvalTask(
                    in_path = fpath,
                    ref_path = ref_path,
                    info = info,
                    structure = info['structure'],
                    name = info['name'],
                    method = info['method'],
                    cdr = info['tag'],
                    ab_chains = info['antibody_chains'],
                    residue_first = info.get('residue_first', None),
                    residue_last  = info.get('residue_last', None),
                ))
                self.visited.add(fpath)
        return tasks

AbDock/src/tools/eval/base.py
pengzhangzhi/ab_opt
# pyright: reportMissingImports=false
import pyrosetta
from pyrosetta.rosetta.protocols.analysis import InterfaceAnalyzerMover
pyrosetta.init(' '.join([
    '-mute', 'all',
    '-use_input_sc',
    '-ignore_unrecognized_res',
    '-ignore_zero_occupancy', 'false',
    '-load_PDB_components', 'false',
    '-relax:default_repeats', '2',
    '-no_fconfig',
]))

from src.tools.eval.base import EvalTask


def pyrosetta_interface_energy(pdb_path, interface):
    pose = pyrosetta.pose_from_pdb(pdb_path)
    mover = InterfaceAnalyzerMover(interface)
    mover.set_pack_separated(True)
    mover.apply(pose)
    return pose.scores['dG_separated']


def eval_interface_energy(task: EvalTask):
    model_gen = task.get_gen_biopython_model()
    antigen_chains = set()
    for chain in model_gen:
        if chain.id not in task.ab_chains:
            antigen_chains.add(chain.id)
    antigen_chains = ''.join(list(antigen_chains))
    antibody_chains = ''.join(task.ab_chains)
    interface = f"{antibody_chains}_{antigen_chains}"

    dG_gen = pyrosetta_interface_energy(task.in_path, interface)
    dG_ref = pyrosetta_interface_energy(task.ref_path, interface)
    print(f"Energy: {dG_gen:.2f} vs {dG_ref:.2f} ({dG_gen - dG_ref:.2f})")
    task.scores.update({
        'dG_gen': dG_gen,
        'dG_ref': dG_ref,
        'ddG': dG_gen - dG_ref
    })
    return task

AbDock/src/tools/eval/energy.py
pengzhangzhi/ab_opt
import os
import argparse
import ray
import shelve
import time
import pandas as pd
from typing import Mapping

from src.tools.eval.base import EvalTask, TaskScanner
from src.tools.eval.similarity import eval_similarity
from src.tools.eval.energy import eval_interface_energy


@ray.remote(num_cpus=1)
def evaluate(task, args):
    return evaluate_local(task, args.no_energy)


def evaluate_local(task, no_energy):
    funcs = []
    funcs.append(eval_similarity)
    if not no_energy:
        funcs.append(eval_interface_energy)
    for f in funcs:
        task = f(task)
    return task

def dump_db(db: Mapping[str, EvalTask], path):
    table = []
    for task in db.values():
        if 'abopt' in path and task.scores['seqid'] >= 100.0:
            # In abopt (Antibody Optimization) mode, ignore sequences identical to the wild-type
            continue
        table.append(task.to_report_dict())
    table = pd.DataFrame(table)
    table.to_csv(path, index=False, float_format='%.6f')
    return table

def run_energy(root, pfx='rosetta', no_energy=False):
    db_path = os.path.join(root, 'evaluation_db')
    with shelve.open(db_path) as db:
        scanner = TaskScanner(root=root, postfix=pfx, db=db)
        tasks = scanner.scan()
        for task in tasks:
            done_task = evaluate_local(task, no_energy)
            done_task.save_to_db(db)
        db.sync()
        dump_db(db, os.path.join(root, 'summary.csv'))
        
def main(args):
    
    db_path = os.path.join(args.root, 'evaluation_db')
    with shelve.open(db_path) as db:
        scanner = TaskScanner(root=args.root, postfix=args.pfx, db=db)

        tasks = scanner.scan()
        futures = [evaluate.remote(t, args) for t in tasks]
        if len(futures) > 0:
            print(f'Submitted {len(futures)} tasks.')
        ray.get(futures)
        # while len(futures) > 0: 
            # done_ids, futures = ray.wait(futures, num_returns=1)
            # for done_id in done_ids:
            #     done_task = ray.get(done_id)
            #     done_task.save_to_db(db)
            #     print(f'Remaining {len(futures)}. Finished {done_task.in_path}')
        db.sync()
        
        dump_db(db, os.path.join(args.root, 'summary.csv'))

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, default='./results')
    parser.add_argument('--pfx', type=str, default='rosetta')
    parser.add_argument('--no_energy', action='store_true', default=False)
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    ray.init()
    args = parse_args()
    main(args)

AbDock/src/tools/eval/run.py
pengzhangzhi/ab_opt
import numpy as np
from Bio.PDB import PDBParser, Selection
from Bio.PDB.Polypeptide import three_to_one
from Bio import pairwise2
from Bio.Align import substitution_matrices

from src.tools.eval.base import EvalTask


def reslist_rmsd(res_list1, res_list2):
    res_short, res_long = (res_list1, res_list2) if len(res_list1) < len(res_list2) else (res_list2, res_list1)
    M, N = len(res_short), len(res_long)

    def d(i, j):
        coord_i = np.array(res_short[i]['CA'].get_coord())
        coord_j = np.array(res_long[j]['CA'].get_coord())
        return ((coord_i - coord_j) ** 2).sum()

    SD = np.full([M, N], np.inf)
    for i in range(M):
        j = N - (M - i)
        SD[i, j] = sum([ d(i+k, j+k) for k in range(N-j) ])
    
    for j in range(N):
        SD[M-1, j] = d(M-1, j)

    for i in range(M-2, -1, -1):
        for j in range((N-(M-i))-1, -1, -1):
            SD[i, j] = min(
                d(i, j) + SD[i+1, j+1],
                SD[i, j+1]
            )

    min_SD = SD[0, :N-M+1].min()
    best_RMSD = np.sqrt(min_SD / M)
    return best_RMSD


def entity_to_seq(entity):
    seq = ''
    mapping = []
    for res in Selection.unfold_entities(entity, 'R'):
        try:
            seq += three_to_one(res.get_resname())
            mapping.append(res.get_id())
        except KeyError:
            pass
    assert len(seq) == len(mapping)
    return seq, mapping


def reslist_seqid(res_list1, res_list2):
    seq1, _ = entity_to_seq(res_list1)
    seq2, _ = entity_to_seq(res_list2)
    _, seq_id = align_sequences(seq1, seq2)
    return seq_id


def align_sequences(sequence_A, sequence_B, **kwargs):
    """
    Performs a global pairwise alignment between two sequences
    using the BLOSUM62 matrix and the Needleman-Wunsch algorithm
    as implemented in Biopython. Returns the alignment, the sequence
    identity and the residue mapping between both original sequences.
    """

    def _calculate_identity(sequenceA, sequenceB):
        """
        Returns the percentage of identical characters between two sequences.
        Assumes the sequences are aligned.
        """

        sa, sb, sl = sequenceA, sequenceB, len(sequenceA)
        matches = [sa[i] == sb[i] for i in range(sl)]
        seq_id = (100 * sum(matches)) / sl
        return seq_id

        # gapless_sl = sum([1 for i in range(sl) if (sa[i] != '-' and sb[i] != '-')])
        # gap_id = (100 * sum(matches)) / gapless_sl
        # return (seq_id, gap_id)

    #
    matrix = kwargs.get('matrix', substitution_matrices.load("BLOSUM62"))
    gap_open = kwargs.get('gap_open', -10.0)
    gap_extend = kwargs.get('gap_extend', -0.5)

    alns = pairwise2.align.globalds(sequence_A, sequence_B,
                                    matrix, gap_open, gap_extend,
                                    penalize_end_gaps=(False, False) )

    best_aln = alns[0]
    aligned_A, aligned_B, score, begin, end = best_aln

    # Calculate sequence identity
    seq_id = _calculate_identity(aligned_A, aligned_B)
    return (aligned_A, aligned_B), seq_id


def extract_reslist(model, residue_first, residue_last):
    assert residue_first[0] == residue_last[0]
    residue_first, residue_last = tuple(residue_first), tuple(residue_last)

    chain_id = residue_first[0]
    pos_first, pos_last = residue_first[1:], residue_last[1:]
    chain = model[chain_id]
    reslist = []
    for res in Selection.unfold_entities(chain, 'R'):
        pos_current = (res.id[1], res.id[2])
        if pos_first <= pos_current <= pos_last:
            reslist.append(res)
    return reslist


def eval_similarity(task: EvalTask):
    model_gen = task.get_gen_biopython_model()
    model_ref = task.get_ref_biopython_model()

    reslist_gen = extract_reslist(model_gen, task.residue_first, task.residue_last)
    reslist_ref = extract_reslist(model_ref, task.residue_first, task.residue_last)

    task.scores.update({
        'rmsd': reslist_rmsd(reslist_gen, reslist_ref),
        'seqid': reslist_seqid(reslist_gen, reslist_ref),
    })
    return task

AbDock/src/tools/eval/similarity.py
pengzhangzhi/ab_opt
import os
import re
import json
from typing import Optional, Tuple, List
from dataclasses import dataclass


@dataclass
class RelaxTask:
    in_path: str
    current_path: str
    info: dict
    status: str

    flexible_residue_first: Optional[Tuple] = None
    flexible_residue_last: Optional[Tuple] = None

    def get_in_path_with_tag(self, tag):
        name, ext = os.path.splitext(self.in_path)
        new_path = f'{name}_{tag}{ext}'
        return new_path

    def set_current_path_tag(self, tag):
        new_path = self.get_in_path_with_tag(tag)
        self.current_path = new_path
        return new_path

    def check_current_path_exists(self):
        ok = os.path.exists(self.current_path)
        if not ok:
            self.mark_failure()
        if os.path.getsize(self.current_path) == 0:
            ok = False
            self.mark_failure()
            os.unlink(self.current_path)
        return ok

    def update_if_finished(self, tag):
        out_path = self.get_in_path_with_tag(tag)
        if os.path.exists(out_path) and os.path.getsize(out_path) > 0:
            print('Already finished', out_path)
            self.set_current_path_tag(tag)
            self.mark_success()
            return True
        return False

    def can_proceed(self):
        self.check_current_path_exists()
        return self.status != 'failed'

    def mark_success(self):
        self.status = 'success'

    def mark_failure(self):
        self.status = 'failed'

    def __str__(self):
        return f'{self.in_path} -> {self.current_path}'
    def __repr__(self):
        return str(self)
class TaskScanner:

    def __init__(self, root, final_postfix=None):
        super().__init__()
        self.root = root
        self.visited = set()
        self.final_postfix = final_postfix

    def _get_metadata(self, fpath):
        json_path = os.path.join(
            os.path.dirname(os.path.dirname(fpath)), 
            'metadata.json'
        )
        tag_name = os.path.basename(os.path.dirname(fpath))
        try:
            with open(json_path, 'r') as f:
                metadata = json.load(f)
            for item in metadata['items']:
                if item['tag'] == tag_name:
                    return item
        except (json.JSONDecodeError, FileNotFoundError) as e:
            return None
        return None

    def scan(self) -> List[RelaxTask]: 
        tasks = []
        input_fname_pattern = '^(?!.*patch)(^\d+\.pdb$|^REF\d\.pdb$)'
        for parent, _, files in os.walk(self.root):
            for fname in files:
                fpath = os.path.join(parent, fname)
                if not re.match(input_fname_pattern, fname):
                    continue
                if os.path.getsize(fpath) == 0:
                    continue
                if fpath in self.visited:
                    continue
                
                # If finished
                if self.final_postfix is not None:
                    fpath_name, fpath_ext = os.path.splitext(fpath)
                    fpath_final = f"{fpath_name}_{self.final_postfix}{fpath_ext}"
                    if os.path.exists(fpath_final):
                        continue

                # Get metadata
                info = self._get_metadata(fpath)
                if info is None:
                    continue
                    
                tasks.append(RelaxTask(
                    in_path = fpath,
                    current_path = fpath,
                    info = info,
                    status = 'created',
                    flexible_residue_first = info.get('residue_first', None),
                    flexible_residue_last  = info.get('residue_last', None),
                ))
                self.visited.add(fpath)
        return tasks

AbDock/src/tools/relax/base.py
pengzhangzhi/ab_opt
import os
import time
import io
import logging
import pdbfixer
import openmm
from openmm import app as openmm_app
from openmm import unit
ENERGY = unit.kilocalories_per_mole
LENGTH = unit.angstroms

from src.tools.relax.base import RelaxTask


def current_milli_time():
    return round(time.time() * 1000)


def _is_in_the_range(ch_rs_ic, flexible_residue_first, flexible_residue_last):
    if ch_rs_ic[0] != flexible_residue_first[0]: return False
    r_first, r_last = tuple(flexible_residue_first[1:]), tuple(flexible_residue_last[1:])
    rs_ic = ch_rs_ic[1:]
    return r_first <= rs_ic <= r_last


class ForceFieldMinimizer(object):

    def __init__(self, stiffness=10.0, max_iterations=0, tolerance=2.39*unit.kilocalories_per_mole, platform='CUDA'):
        super().__init__()
        self.stiffness = stiffness
        self.max_iterations = max_iterations
        self.tolerance = tolerance
        assert platform in ('CUDA', 'CPU')
        self.platform = platform

    def _fix(self, pdb_str):
        fixer = pdbfixer.PDBFixer(pdbfile=io.StringIO(pdb_str))
        fixer.findNonstandardResidues()
        fixer.replaceNonstandardResidues()

        fixer.findMissingResidues()
        fixer.findMissingAtoms()
        fixer.addMissingAtoms(seed=0)
        fixer.addMissingHydrogens()

        out_handle = io.StringIO()
        openmm_app.PDBFile.writeFile(fixer.topology, fixer.positions, out_handle, keepIds=True)
        return out_handle.getvalue()

    def _get_pdb_string(self, topology, positions):
        with io.StringIO() as f:
            openmm_app.PDBFile.writeFile(topology, positions, f, keepIds=True)
            return f.getvalue()

    def _minimize(self, pdb_str, flexible_residue_first=None, flexible_residue_last=None):
        pdb = openmm_app.PDBFile(io.StringIO(pdb_str))

        force_field = openmm_app.ForceField("amber99sb.xml")
        constraints = openmm_app.HBonds
        system = force_field.createSystem(pdb.topology, constraints=constraints)

        # Add constraints to non-generated regions
        force = openmm.CustomExternalForce("0.5 * k * ((x-x0)^2 + (y-y0)^2 + (z-z0)^2)")
        force.addGlobalParameter("k", self.stiffness)
        for p in ["x0", "y0", "z0"]:
            force.addPerParticleParameter(p)
        
        if flexible_residue_first is not None and flexible_residue_last is not None:
            for i, a in enumerate(pdb.topology.atoms()):
                ch_rs_ic = (a.residue.chain.id, int(a.residue.id), a.residue.insertionCode)
                if not _is_in_the_range(ch_rs_ic, flexible_residue_first, flexible_residue_last) and a.element.name != "hydrogen":
                    force.addParticle(i, pdb.positions[i])
                
        system.addForce(force)

        # Set up the integrator and simulation
        integrator = openmm.LangevinIntegrator(0, 0.01, 0.0)
        platform = openmm.Platform.getPlatformByName("CUDA")
        simulation = openmm_app.Simulation(pdb.topology, system, integrator, platform)
        simulation.context.setPositions(pdb.positions)

        # Perform minimization
        ret = {}
        state = simulation.context.getState(getEnergy=True, getPositions=True)
        ret["einit"] = state.getPotentialEnergy().value_in_unit(ENERGY)
        ret["posinit"] = state.getPositions(asNumpy=True).value_in_unit(LENGTH)

        simulation.minimizeEnergy(maxIterations=self.max_iterations, tolerance=self.tolerance)

        state = simulation.context.getState(getEnergy=True, getPositions=True)
        ret["efinal"] = state.getPotentialEnergy().value_in_unit(ENERGY)
        ret["pos"] = state.getPositions(asNumpy=True).value_in_unit(LENGTH)
        ret["min_pdb"] = self._get_pdb_string(simulation.topology, state.getPositions())

        return ret['min_pdb'], ret

    def _add_energy_remarks(self, pdb_str, ret):
        pdb_lines = pdb_str.splitlines()
        pdb_lines.insert(1, "REMARK   1  FINAL ENERGY:   {:.3f} KCAL/MOL".format(ret['efinal']))
        pdb_lines.insert(1, "REMARK   1  INITIAL ENERGY: {:.3f} KCAL/MOL".format(ret['einit']))
        return "\n".join(pdb_lines)

    def __call__(self, pdb_str, flexible_residue_first=None, flexible_residue_last=None, return_info=True):
        if '\n' not in pdb_str and pdb_str.lower().endswith(".pdb"):
            with open(pdb_str) as f:
                pdb_str = f.read()

        pdb_fixed = self._fix(pdb_str)
        pdb_min, ret = self._minimize(pdb_fixed, flexible_residue_first, flexible_residue_last)
        pdb_min = self._add_energy_remarks(pdb_min, ret)
        if return_info:
            return pdb_min, ret
        else:
            return pdb_min


def run_openmm(task: RelaxTask):
    if not task.can_proceed() :
        return task
    if task.update_if_finished('openmm'):
        return task

    try:
        minimizer = ForceFieldMinimizer()
        with open(task.current_path, 'r') as f:
            pdb_str = f.read()

        pdb_min = minimizer(
            pdb_str = pdb_str,
            flexible_residue_first = task.flexible_residue_first,
            flexible_residue_last = task.flexible_residue_last,
            return_info = False,
        )
        out_path = task.set_current_path_tag('openmm')
        with open(out_path, 'w') as f:
            f.write(pdb_min)
        task.mark_success()
    except ValueError as e:
        logging.warning(
            f'{e.__class__.__name__}: {str(e)} ({task.current_path})'
        )        
        task.mark_failure()
    return task


AbDock/src/tools/relax/openmm_relaxer.py
pengzhangzhi/ab_opt
# pyright: reportMissingImports=false
import os
import time
import pyrosetta
from pyrosetta.rosetta.protocols.relax import FastRelax
from pyrosetta.rosetta.core.pack.task import TaskFactory
from pyrosetta.rosetta.core.pack.task import operation
from pyrosetta.rosetta.core.select import residue_selector as selections
from pyrosetta.rosetta.core.select.movemap import MoveMapFactory, move_map_action
pyrosetta.init(' '.join([
    '-mute', 'all',
    '-use_input_sc',
    '-ignore_unrecognized_res',
    '-ignore_zero_occupancy', 'false',
    '-load_PDB_components', 'false',
    '-relax:default_repeats', '2',
    '-no_fconfig',
]))

from src.tools.relax.base import RelaxTask


def current_milli_time():
    return round(time.time() * 1000)


def parse_residue_position(p):
    icode = None
    if not p[-1].isnumeric():   # Has ICODE
        icode = p[-1]

    for i, c in enumerate(p):
        if c.isnumeric():
            break
    chain = p[:i]
    resseq = int(p[i:])

    if icode is not None:
        return chain, resseq, icode
    else:
        return chain, resseq


def get_scorefxn(scorefxn_name:str):
    """
    Gets the scorefxn with appropriate corrections.
    Taken from: https://gist.github.com/matteoferla/b33585f3aeab58b8424581279e032550
    """
    import pyrosetta

    corrections = {
        'beta_july15': False,
        'beta_nov16': False,
        'gen_potential': False,
        'restore_talaris_behavior': False,
    }
    if 'beta_july15' in scorefxn_name or 'beta_nov15' in scorefxn_name:
        # beta_july15 is ref2015
        corrections['beta_july15'] = True
    elif 'beta_nov16' in scorefxn_name:
        corrections['beta_nov16'] = True
    elif 'genpot' in scorefxn_name:
        corrections['gen_potential'] = True
        pyrosetta.rosetta.basic.options.set_boolean_option('corrections:beta_july15', True)
    elif 'talaris' in scorefxn_name:  #2013 and 2014
        corrections['restore_talaris_behavior'] = True
    else:
        pass
    for corr, value in corrections.items():
        pyrosetta.rosetta.basic.options.set_boolean_option(f'corrections:{corr}', value)
    return pyrosetta.create_score_function(scorefxn_name)


class RelaxRegion(object):
    
    def __init__(self, scorefxn='ref2015', max_iter=1000, subset='nbrs', move_bb=True):
        super().__init__()
        self.scorefxn = get_scorefxn(scorefxn)
        self.fast_relax = FastRelax()
        self.fast_relax.set_scorefxn(self.scorefxn)
        self.fast_relax.max_iter(max_iter)
        assert subset in ('all', 'target', 'nbrs')
        self.subset = subset
        self.move_bb = move_bb

    def __call__(self, pdb_path, flexible_residue_first, flexible_residue_last):
        pose = pyrosetta.pose_from_pdb(pdb_path)
        start_t = current_milli_time()
        original_pose = pose.clone()

        tf = TaskFactory()
        tf.push_back(operation.InitializeFromCommandline())
        tf.push_back(operation.RestrictToRepacking())   # Only allow residues to repack. No design at any position.

        # Create selector for the region to be relaxed
        # Turn off design and repacking on irrelevant positions
        if flexible_residue_first[-1] == ' ': 
            flexible_residue_first = flexible_residue_first[:-1]
        if flexible_residue_last[-1] == ' ':  
            flexible_residue_last  = flexible_residue_last[:-1]
        if self.subset != 'all':
            gen_selector = selections.ResidueIndexSelector()
            gen_selector.set_index_range(
                pose.pdb_info().pdb2pose(*flexible_residue_first), 
                pose.pdb_info().pdb2pose(*flexible_residue_last), 
            )
            nbr_selector = selections.NeighborhoodResidueSelector()
            nbr_selector.set_focus_selector(gen_selector)
            nbr_selector.set_include_focus_in_subset(True)

            if self.subset == 'nbrs':
                subset_selector = nbr_selector
            elif self.subset == 'target':
                subset_selector = gen_selector

            prevent_repacking_rlt = operation.PreventRepackingRLT()
            prevent_subset_repacking = operation.OperateOnResidueSubset(
                prevent_repacking_rlt, 
                subset_selector,
                flip_subset=True,
            )
            tf.push_back(prevent_subset_repacking)

        scorefxn = self.scorefxn
        fr = self.fast_relax

        pose = original_pose.clone()
        pos_list = pyrosetta.rosetta.utility.vector1_unsigned_long()
        for pos in range(pose.pdb_info().pdb2pose(*flexible_residue_first), pose.pdb_info().pdb2pose(*flexible_residue_last)+1):
            pos_list.append(pos)
        # basic_idealize(pose, pos_list, scorefxn, fast=True)

        mmf = MoveMapFactory()
        if self.move_bb: 
            mmf.add_bb_action(move_map_action.mm_enable, gen_selector)
        mmf.add_chi_action(move_map_action.mm_enable, subset_selector)
        mm  = mmf.create_movemap_from_pose(pose)

        fr.set_movemap(mm)
        fr.set_task_factory(tf)
        fr.apply(pose)

        e_before = scorefxn(original_pose)
        e_relax  = scorefxn(pose) 
        # print('\n\n[Finished in %.2f secs]' % ((current_milli_time() - start_t) / 1000))
        # print(' > Energy (before):    %.4f' % scorefxn(original_pose))
        # print(' > Energy (optimized): %.4f' % scorefxn(pose))
        return pose, e_before, e_relax


def run_pyrosetta(task: RelaxTask):
    if not task.can_proceed() :
        return task
    if task.update_if_finished('rosetta'):
        return task

    minimizer = RelaxRegion()
    pose_min, _, _ = minimizer(
        pdb_path = task.current_path,
        flexible_residue_first = task.flexible_residue_first,
        flexible_residue_last = task.flexible_residue_last,
    )

    out_path = task.set_current_path_tag('rosetta')
    pose_min.dump_pdb(out_path)
    task.mark_success()
    return task


def run_pyrosetta_fixbb(task: RelaxTask):
    if not task.can_proceed() :
        return task
    if task.update_if_finished('fixbb'):
        return task

    minimizer = RelaxRegion(move_bb=False)
    pose_min, _, _ = minimizer(
        pdb_path = task.current_path,
        flexible_residue_first = task.flexible_residue_first,
        flexible_residue_last = task.flexible_residue_last,
    )

    out_path = task.set_current_path_tag('fixbb')
    pose_min.dump_pdb(out_path)
    task.mark_success()
    return task

    


AbDock/src/tools/relax/pyrosetta_relaxer.py
pengzhangzhi/ab_opt
import argparse
import ray
import time

from src.tools.relax.openmm_relaxer import run_openmm
from src.tools.relax.pyrosetta_relaxer import run_pyrosetta, run_pyrosetta_fixbb
from src.tools.relax.base import TaskScanner


@ray.remote(num_gpus=1/8, num_cpus=1)
def run_openmm_remote(task):
    return run_openmm(task)

def _run_openmm(task):
    return run_openmm(task)

@ray.remote(num_cpus=1)
def run_pyrosetta_remote(task):
    return run_pyrosetta(task)

def _run_pyrosetta(task):
    return run_pyrosetta(task)


@ray.remote(num_cpus=1)
def run_pyrosetta_fixbb_remote(task):
    return run_pyrosetta_fixbb(task)

def _run_pyrosetta_fixbb(task):
    return run_pyrosetta_fixbb(task)

@ray.remote
def pipeline_openmm_pyrosetta(task):
    funcs = [
        run_openmm_remote,
        run_pyrosetta_remote,
    ]
    for fn in funcs:
        task = fn.remote(task)
    return ray.get(task)

def _pipeline_openmm_pyrosetta(task):
    funcs = [
        _run_openmm,
        _run_pyrosetta,
    ]
    for fn in funcs:
        task = fn(task)
    return task

@ray.remote
def pipeline_pyrosetta(task):
    funcs = [
        run_pyrosetta_remote,
    ]
    for fn in funcs:
        task = fn.remote(task)
    return ray.get(task)

def _pipeline_pyrosetta(task):
    funcs = [
        _run_pyrosetta,
    ]
    for fn in funcs:
        task = fn(task)
    return task

@ray.remote
def pipeline_pyrosetta_fixbb(task):
    funcs = [
        run_pyrosetta_fixbb_remote,
    ]
    for fn in funcs:
        task = fn.remote(task)
    return ray.get(task)

def _pipeline_pyrosetta_fixbb(task):
    funcs = [
        _run_pyrosetta_fixbb,
    ]
    for fn in funcs:
        task = fn(task)
    return task


pipeline_dict = {
    'openmm_pyrosetta': pipeline_openmm_pyrosetta,
    'pyrosetta': pipeline_pyrosetta,
    'pyrosetta_fixbb': pipeline_pyrosetta_fixbb,
}

pipeline_dict_local = {
    'openmm_pyrosetta': _pipeline_openmm_pyrosetta,
    'pyrosetta': _pipeline_pyrosetta,
    'pyrosetta_fixbb': _pipeline_pyrosetta_fixbb,
}
def run_relax(root, pipeline_name='openmm_pyrosetta'):
    assert pipeline_name in pipeline_dict_local.keys()
    pipeline = pipeline_dict_local[pipeline_name]
    final_pfx = 'fixbb' if pipeline == pipeline_pyrosetta_fixbb else 'rosetta'
    scanner = TaskScanner(root, final_postfix=final_pfx)
    tasks = scanner.scan()
    for t in tasks:
        pipeline(t)
        
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', type=str, default='./results')
    parser.add_argument('--pipeline', type=lambda s: pipeline_dict[s], default=pipeline_openmm_pyrosetta)
    args = parser.parse_args()
    return args


def main(args):
    final_pfx = 'fixbb' if args.pipeline == pipeline_pyrosetta_fixbb else 'rosetta'
    scanner = TaskScanner(args.root, final_postfix=final_pfx)
    tasks = scanner.scan()
    futures = [args.pipeline.remote(t) for t in tasks]
    if len(futures) > 0:
        print(f'Submitted {len(futures)} tasks.')
    ray.get(futures)
    # while len(futures) > 0:
        # done_ids, futures = ray.wait(futures, num_returns=1)
        # for done_id in done_ids:
        #     done_task = ray.get(done_id)
        #     print(f'Remaining {len(futures)}. Finished {done_task.current_path}')

if __name__ == '__main__':
    ray.init()
    args = parse_args()
    main(args)

AbDock/src/tools/relax/run.py
pengzhangzhi/ab_opt
from .run import renumber

AbDock/src/tools/renumber/__init__.py
pengzhangzhi/ab_opt
import argparse
import abnumber
from Bio import PDB
from Bio.PDB import Model, Chain, Residue, Selection
from Bio.Data import SCOPData
from typing import List, Tuple


def biopython_chain_to_sequence(chain: Chain.Chain):
    residue_list = Selection.unfold_entities(chain, 'R')
    seq = ''.join([SCOPData.protein_letters_3to1.get(r.resname, 'X') for r in residue_list])
    return seq, residue_list


def assign_number_to_sequence(seq, scheme='chothia'):
    abchain = abnumber.Chain(seq, scheme=scheme)
    offset = seq.index(abchain.seq)
    if not (offset >= 0):
        raise ValueError(
            'The identified Fv sequence is not a subsequence of the original sequence.'
        )

    numbers = [None for _ in range(len(seq))]
    for i, (pos, aa) in enumerate(abchain):
        resseq = pos.number
        icode = pos.letter if pos.letter else ' '
        numbers[i+offset] = (resseq, icode)
    return numbers, abchain


def renumber_biopython_chain(chain_id, residue_list: List[Residue.Residue], numbers: List[Tuple[int, str]]):
    chain = Chain.Chain(chain_id)
    for residue, number in zip(residue_list, numbers):
        if number is None:
            continue
        residue = residue.copy()
        new_id = (residue.id[0], number[0], number[1])
        residue.id = new_id
        chain.add(residue)
    return chain


def renumber(in_pdb, out_pdb, return_other_chains=False,scheme='chothia'):
    parser = PDB.PDBParser(QUIET=True)
    structure = parser.get_structure(None, in_pdb)
    model = structure[0]
    model_new = Model.Model(0)

    heavy_chains, light_chains, other_chains = [], [], []

    for chain in model:
        try:
            seq, reslist = biopython_chain_to_sequence(chain)
            numbers, abchain = assign_number_to_sequence(seq,scheme)
            chain_new = renumber_biopython_chain(chain.id, reslist, numbers)
            print(f'[INFO] Renumbered chain {chain_new.id} ({abchain.chain_type})')
            if abchain.chain_type == 'H':
                heavy_chains.append(chain_new.id)
            elif abchain.chain_type in ('K', 'L'):
                light_chains.append(chain_new.id)
        except abnumber.ChainParseError as e:
            print(f'[INFO] Chain {chain.id} does not contain valid Fv: {str(e)}')
            chain_new = chain.copy()
            other_chains.append(chain_new.id)
        model_new.add(chain_new)

    pdb_io = PDB.PDBIO()
    pdb_io.set_structure(model_new)
    pdb_io.save(out_pdb)
    if return_other_chains:
        return heavy_chains, light_chains, other_chains
    else:
        return heavy_chains, light_chains


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('in_pdb', type=str)
    parser.add_argument('out_pdb', type=str)
    args = parser.parse_args()

    renumber(args.in_pdb, args.out_pdb)

if __name__ == '__main__':
    main()

AbDock/src/tools/renumber/run.py
pengzhangzhi/ab_opt
from collections import defaultdict
import os
import argparse
import copy
import json
from tqdm.auto import tqdm
from torch.utils.data import DataLoader
from DockQ.DockQ import calc_DockQ
# from src.tools.relax.run import run_relax
from src.utils.transforms._base import _mask_select_data
import pandas as pd
from src.datasets.custom import preprocess_antibody_structure
from src.models import get_model
from src.modules.common.geometry import reconstruct_backbone_partially
from src.modules.common.so3 import so3vec_to_rotation
from src.tools.runner.design_for_testset import (
    create_data_variants,
    extract_dict,
    rank_commoness,
    traverse_dict,
    calc_avg_rmsd,
)
from src.utils.inference import RemoveNative
from src.utils.protein.writers import save_pdb
from src.utils.train import recursive_to
from src.utils.misc import *
from src.utils.data import *
from src.utils.transforms import *
from src.utils.inference import *
from src.tools.renumber import renumber as renumber_antibody


def dock_for_pdb(args):
    # Load configs
    config, config_name = load_config(args.config)
    seed_all(args.seed if args.seed is not None else config.sampling.seed)

    # Structure loading
    data_id = os.path.basename(args.pdb_path) if args.id == '' else args.id
    if args.label_heavy_as_cdr:
        label_whole_heavy_chain_as_cdr = True
        pdb_path = args.pdb_path
        assert (
            args.heavy is not None
        ), f"must specify heavy chain id for seq designed pdb."
    else:
        label_whole_heavy_chain_as_cdr = False
        if args.no_renumber:
            pdb_path = args.pdb_path
        else:
            in_pdb_path = args.pdb_path
            out_pdb_path = os.path.splitext(in_pdb_path)[0] + "_chothia.pdb"
            heavy_chains, light_chains = renumber_antibody(in_pdb_path, out_pdb_path)
            pdb_path = out_pdb_path

            if args.heavy is None and len(heavy_chains) > 0:
                args.heavy = heavy_chains[0]
            if args.light is None and len(light_chains) > 0:
                args.light = light_chains[0]
        if args.heavy is None and args.light is None:
            raise ValueError(
                "Neither heavy chain id (--heavy) or light chain id (--light) is specified."
            )
    get_structure = lambda: preprocess_antibody_structure(
        {
            "id": data_id,
            "pdb_path": pdb_path,
            "heavy_id": args.heavy,
            # If the input is a nanobody, the light chain will be ignores
            "light_id": args.light,
        },
        label_whole_heavy_chain_as_cdr=label_whole_heavy_chain_as_cdr,
    )

    # Logging
    structure_ = get_structure()
    structure_id = structure_["id"]
    tag_postfix = "_%s" % args.tag if args.tag else ""
    log_dir = get_new_log_dir(
        os.path.join(args.out_root, config_name + tag_postfix), prefix=data_id
    )
    logger = get_logger("sample", log_dir)
    logger.info("Data ID: %s" % structure_["id"])
    logger.info(f"Results will be saved to {log_dir}")
    data_native = MergeChains()(structure_)
    save_pdb(data_native, os.path.join(log_dir, "reference.pdb"))

    # Load checkpoint and model
    logger.info("Loading model config and checkpoints: %s" % (args.ckpt))
    ckpt = torch.load(args.ckpt, map_location="cpu")
    cfg_ckpt = ckpt["config"]
    model_cfg = cfg_ckpt.model
    model = get_model(model_cfg).to(args.device)
    lsd = model.load_state_dict(ckpt["model"])
    logger.info(str(lsd))

    # Make data variants
    data_variants = create_data_variants(
        config=config,
        structure_factory=get_structure,
    )

    # Save metadata
    metadata = {
        "identifier": structure_id,
        "config": args.config,
        "items": [
            {kk: vv for kk, vv in var.items() if kk != "data"} for var in data_variants
        ],
    }
    with open(os.path.join(log_dir, "metadata.json"), "w") as f:
        json.dump(metadata, f, indent=2)

    # Start sampling
    collate_fn = PaddingCollate(eight=False)
    inference_tfm = [
        PatchAroundAnchor(
            initial_patch_size=model_cfg.initial_patch_size,
            antigen_size=model_cfg.antigen_size,
            remove_anchor=model_cfg.remove_anchor,
            crop_contiguous_antigen=model_cfg.get("crop_contiguous_antigen", False),
            contiguous_threshold=model_cfg.get("contiguous_threshold", 1e6),
        ),
    ]
    if "abopt" not in config.mode and args.contig == '':  # Don't remove native CDR in optimization mode
        inference_tfm.append(
            RemoveNative(
                remove_structure=config.sampling.sample_structure,
                remove_sequence=config.sampling.sample_sequence,
            )
        )
    inference_tfm = Compose(inference_tfm)
    result_dict = {}
    aa_df = pd.DataFrame(columns=["Region", "native_aa", "sampled_aa", "AAR" , "PPL"])
    for variant in data_variants:
        variant_result_dict = defaultdict(list)
        os.makedirs(os.path.join(log_dir, variant["tag"]), exist_ok=True)
        logger.info(f"Start sampling for: {variant['tag']}")

        data_cropped = inference_tfm(copy.deepcopy(variant["data"]))
        data_list_repeat = [data_cropped] * args.num_samples
        loader = DataLoader(
            data_list_repeat,
            batch_size=args.batch_size,
            shuffle=False,
            collate_fn=collate_fn,
        )

        count = 0
        candidates = []
        for batch in tqdm(loader, desc=variant["name"], dynamic_ncols=True):
            torch.set_grad_enabled(False)
            model.eval()
            batch = recursive_to(batch, args.device)
            traj_batch = model.sample(
                batch,
                sample_opt={
                    "pbar": True,
                    "sample_structure": config.sampling.sample_structure,
                    "sample_sequence": config.sampling.sample_sequence,
                    "contig": args.contig,
                },
            )
            batch = recursive_to(batch, 'cpu')
            traj_batch = recursive_to(traj_batch, 'cpu')
            def save_traj(traj_batch, save_path):
                pdb_str = ""
                for i in reversed(range(len(traj_batch))):
                    aa_new_i = traj_batch[i][2]  
                    pos_atom_new_i, mask_atom_new_i = reconstruct_backbone_partially(
                        pos_ctx=batch["pos_heavyatom"],
                        R_new=so3vec_to_rotation(traj_batch[i][0]),
                        t_new=traj_batch[i][1],
                        aa=aa_new_i,
                        chain_nb=batch["chain_nb"],
                        res_nb=batch["res_nb"],
                        mask_atoms=batch["mask_heavyatom"],
                        mask_recons=batch["generate_flag"],
                    )   
                    aa_new_i = aa_new_i.cpu()[0]
                    pos_atom_new_i = pos_atom_new_i.cpu()[0]
                    mask_atom_new_i = mask_atom_new_i.cpu()[0]
                    data_tmpl = variant["data"]
                    antigen_mask = data_tmpl["fragment_type"] == constants.Fragment.Antigen
                    patch_mask = torch.zeros_like(antigen_mask).bool()
                    patch_mask[data_cropped["patch_idx"]] = True
                    antigen_and_patch_mask = antigen_mask | patch_mask
                    data = copy.deepcopy(data_tmpl)
                    mask_ha = apply_patch_to_tensor(
                        data_tmpl["mask_heavyatom"],
                        mask_atom_new_i,
                        data_cropped["patch_idx"],
                    )
                    pos_ha = apply_patch_to_tensor(
                        data_tmpl["pos_heavyatom"],
                        pos_atom_new_i + batch["origin"][0].view(1, 1, 3).cpu(),
                        data_cropped["patch_idx"],
                    )
                    data["mask_heavyatom"] = mask_ha
                    data["pos_heavyatom"] = pos_ha
                    data_patch = _mask_select_data(data, antigen_and_patch_mask)
                    save_pdb(data_patch, path='tmp.pdb')
                    with open('tmp.pdb', 'r') as f:
                        lines = f.read().splitlines()
                        lines = [f'MODEL     {i+1}'] + lines + ['ENDMDL\n']
                        pdb_str += '\n'.join(lines)
                with open(save_path, 'w') as f:
                    f.write(pdb_str)
            save_traj(traj_batch, os.path.join(log_dir, "traj.pdb"))
            
            aa_new = traj_batch[0][2]  # 0: Last sampling step. 2: Amino acid.
            prmsd = traj_batch[0][3]  # 3: prmsd, [B, N]
            perplexity = traj_batch[0][4]  # 4: perplexity, [B,]
            pos_atom_new, mask_atom_new = reconstruct_backbone_partially(
                pos_ctx=batch["pos_heavyatom"],
                R_new=so3vec_to_rotation(traj_batch[0][0]),
                t_new=traj_batch[0][1],
                aa=aa_new,
                chain_nb=batch["chain_nb"],
                res_nb=batch["res_nb"],
                mask_atoms=batch["mask_heavyatom"],
                mask_recons=batch["generate_flag"],
            )
            aa_new = aa_new.cpu()
            pos_atom_new = pos_atom_new.cpu()
            mask_atom_new = mask_atom_new.cpu()

            # save ref docking pose
            data_tmpl = variant["data"]
            antigen_mask = data_tmpl["fragment_type"] == constants.Fragment.Antigen
            patch_mask = torch.zeros_like(antigen_mask).bool()
            patch_mask[data_cropped["patch_idx"]] = True
            antigen_and_patch_mask = antigen_mask | patch_mask
            native_patch = _mask_select_data(data_tmpl, antigen_and_patch_mask)
            ref_path = os.path.join(log_dir, variant["tag"], "REF1.pdb")
            save_pdb(native_patch, path=ref_path)

            for i in range(aa_new.size(0)):
                data_tmpl = variant["data"]
                generate_flag_i = batch["generate_flag"][i].cpu()
                prmsd_i = round(
                    float(prmsd[i].mean()), 6
                )  # round(float(prmsd[i][generate_flag_i].mean()),3)
                perplexity_i = round(float(perplexity[i]), 6)
                pred_aa = aa_new[i][generate_flag_i]
                aa = apply_patch_to_tensor(
                    data_tmpl["aa"], aa_new[i], data_cropped["patch_idx"]
                )
                mask_ha = apply_patch_to_tensor(
                    data_tmpl["mask_heavyatom"],
                    mask_atom_new[i],
                    data_cropped["patch_idx"],
                )
                pos_ha = apply_patch_to_tensor(
                    data_tmpl["pos_heavyatom"],
                    pos_atom_new[i] + batch["origin"][i].view(1, 1, 3).cpu(),
                    data_cropped["patch_idx"],
                )
                gen_flag_full = apply_patch_to_tensor(
                    data_tmpl["generate_flag"], generate_flag_i, data_cropped["patch_idx"]
                )
                native_aa = data_tmpl["aa"][gen_flag_full]
                generated_structure = pos_ha[gen_flag_full]
                candidates.append(generated_structure)
                aar = float((native_aa == pred_aa).sum().float() / native_aa.size(0))
                native_aa_codes = "".join(
                    constants.aaidx2symbol[int(i)] for i in native_aa
                )
                pred_aa_codes = "".join(constants.aaidx2symbol[int(i)] for i in pred_aa)
                aa_df = pd.concat(
                    [
                        aa_df,
                        pd.DataFrame(
                            {
                                "Region": [variant["tag"]],
                                "native_aa": [native_aa_codes],
                                "sampled_aa": [pred_aa_codes],
                                "AAR": [aar],
                                "PPL": [perplexity_i],
                            }
                        ),
                    ]
                )
                # save cdr + antigen complex.
                antigen_mask = data_tmpl["fragment_type"] == constants.Fragment.Antigen
                patch_mask = torch.zeros_like(antigen_mask).bool()
                patch_mask[data_cropped["patch_idx"]] = True
                antigen_and_patch_mask = antigen_mask | patch_mask
                data = copy.deepcopy(data_tmpl)
                data["aa"] = aa
                data["mask_heavyatom"] = mask_ha
                data["pos_heavyatom"] = pos_ha
                data_patch = _mask_select_data(data, antigen_and_patch_mask)
                save_path = os.path.join(log_dir, variant["tag"], "%04d.pdb" % (count,))
                save_pdb(data_patch, path=save_path)
                del data, data_patch
                
                patch_path = os.path.join(
                    log_dir, variant["tag"], "%04d_patch.pdb" % (count,)
                )
                save_pdb(
                    {
                        "chain_nb": data_cropped["chain_nb"],
                        "chain_id": data_cropped["chain_id"],
                        "resseq": data_cropped["resseq"],
                        "icode": data_cropped["icode"],
                        # Generated
                        "aa": aa_new[i],
                        "mask_heavyatom": mask_atom_new[i],
                        "pos_heavyatom": pos_atom_new[i]
                        + batch["origin"][i].view(1, 1, 3).cpu(),
                    },
                    path=patch_path,
                )
                count += 1
                score_dict = {}
                dock_dict = {
                    k: round(v, 3)
                    for k, v in calc_DockQ(save_path, ref_path, use_CA_only=True).items()
                    if k in ["DockQ", "irms", "Lrms", "fnat"]
                }
                score_dict.update(dock_dict)
                score_dict.update({"AAR": aar, "prmsd": prmsd_i})
                [variant_result_dict[key].append(score_dict[key]) for key in score_dict]
            logger.info("Finished.\n")
        candidates = torch.stack(candidates, dim=0)  # (n_candidates, n_res, n_atoms, 3)
        candidates = candidates[
            ...,
            (
                constants.BBHeavyAtom.N,
                constants.BBHeavyAtom.CA,
                constants.BBHeavyAtom.CB,
            ),
            :,
        ]
        candidates = candidates.reshape(count, -1, 3)
        topk_idxs = rank_commoness(candidates, args.topk)
        avg_rmsd = float(calc_avg_rmsd(candidates))
        top_result = {
            f"{k}_top{args.topk}": [v[i] for i in topk_idxs]
            for k, v in variant_result_dict.items() if isinstance(v, list)
        }
        variant_result_dict.update({"per_sample_rmsd": avg_rmsd})
        
        variant_result_dict.update(top_result)
        result_dict[variant["tag"]] = variant_result_dict
    results = traverse_dict(
        result_dict, list, lambda x: round(float(np.mean(x)), 3), lambda x: f"{x}_mean"
    )
    # mean_results = traverse_dict(result_dict, list, lambda x: round(float(np.mean(x)),3), lambda x: x)
    topk_results = extract_dict(
        result_dict, list, lambda k: k.endswith(f"_top{args.topk}")
    )
    df = pd.DataFrame.from_dict(results, orient="index")
    df.to_csv(os.path.join(log_dir, "results.csv"))
    aa_df.to_csv(os.path.join(log_dir, "aa.csv"), index=False)

    return topk_results  # mean_results


def args_from_cmdline():
 
    parser = argparse.ArgumentParser()
    parser.add_argument("--id", type=str, default="", help="name of the pdb, use for saved dir. if not specified, use the pdb name.")
    parser.add_argument(
        "--pdb_path",
        type=str,
        default="data/examples/7DK2_AB_C.pdb",
        help="Path to the PDB file.",
    )
    parser.add_argument(
        "--label_heavy_as_cdr",
        default=False,
        action="store_true",
        help="Label the heavy chain as CDR. \
            If True, specify the heavy chain id with --heavy. \
        Usually used when the heavy chain contains only designed CDRH3.",
    )
    parser.add_argument(
        "--contig",
        default='',
        help="Specify contiguous residues to be designed. \
            e.g., 1-10, meaning the first to tenth CDR residues will be designed. \
            If not specified, all residues will be designed.",
    )

    parser.add_argument(
        "-c", "--config", type=str, default="configs/test/dock_cdr.yml"
    )
    parser.add_argument(
        "-ck", "--ckpt", type=str, default="reproduction/dock_single_cdr/250000.pt"
    )
    parser.add_argument(
        "--heavy", type=str, default=None, help="Chain id of the heavy chain."
    )
    parser.add_argument(
        "--light", type=str, default=None, help="Chain id of the light chain."
    )
    parser.add_argument("-n", "--num_samples", type=int, default=10)
    parser.add_argument(
        "--topk",
        type=int,
        default=1,
        help="Select topk samples. must be less than or equal to num_samples.",
    )
    parser.add_argument("--no_renumber", action="store_true", default=False)
    parser.add_argument("-o", "--out_root", type=str, default="./results/")
    parser.add_argument("-t", "--tag", type=str, default="")
    parser.add_argument("-s", "--seed", type=int, default=None)
    parser.add_argument("-d", "--device", type=str, default="cuda")
    parser.add_argument("-b", "--batch_size", type=int, default=16)
    args = parser.parse_args()
    assert args.topk <= args.num_samples
    return args


def args_factory(**kwargs):
    default_args = EasyDict(
        heavy="H",
        light="L",
        no_renumber=False,
        config="./configs/test/codesign_single.yml",
        out_root="./results",
        tag="",
        seed=None,
        device="cuda",
        batch_size=16,
    )
    default_args.update(kwargs)
    return default_args


if __name__ == "__main__":
    dock_for_pdb(args_from_cmdline())

AbDock/src/tools/runner/design_for_pdb.py
pengzhangzhi/ab_opt
from collections import defaultdict
import os
import argparse
import copy
import json
from tqdm.auto import tqdm
from torch.utils.data import DataLoader
from src.datasets import get_dataset
from src.models import get_model
from src.modules.common.geometry import reconstruct_backbone_partially
from src.modules.common.so3 import so3vec_to_rotation
from src.tools.eval.run import run_energy
from src.tools.relax.run import run_relax
from src.utils.inference import RemoveNative
from src.utils.protein.writers import save_pdb
from src.utils.train import recursive_to
from src.utils.misc import *
from src.utils.data import *
from src.utils.transforms import *
from src.utils.inference import *
from DockQ.DockQ import calc_DockQ
from src.utils.transforms._base import _mask_select_data
from easydict import EasyDict as edict
import pandas as pd


def create_data_variants(config, structure_factory):
    structure = structure_factory()
    structure_id = structure["id"]

    data_variants = []
    if config.mode == "single_cdr":
        cdrs = sorted(
            list(set(find_cdrs(structure)).intersection(config.sampling.cdrs))
        )
        for cdr_name in cdrs:
            transform = Compose(
                [
                    MaskSingleCDR(cdr_name, augmentation=False),
                    MergeChains(),
                ]
            )
            data_var = transform(structure_factory())
            residue_first, residue_last = get_residue_first_last(data_var)
            data_variants.append(
                {
                    "data": data_var,
                    "name": f"{structure_id}-{cdr_name}",
                    "tag": f"{cdr_name}",
                    "cdr": cdr_name,
                    "residue_first": residue_first,
                    "residue_last": residue_last,
                }
            )
    elif config.mode == "multiple_cdrs":
        cdrs = sorted(
            list(set(find_cdrs(structure)).intersection(config.sampling.cdrs))
        )
        transform = Compose(
            [
                MaskMultipleCDRs(selection=cdrs, augmentation=False),
                MergeChains(),
            ]
        )
        data_var = transform(structure_factory())
        data_variants.append(
            {
                "data": data_var,
                "name": f"{structure_id}-MultipleCDRs",
                "tag": "MultipleCDRs",
                "cdrs": cdrs,
                "residue_first": None,
                "residue_last": None,
            }
        )
    elif config.mode == "dock_antibody":
        transform = Compose(
            [
                MaskFullAntibody(antibody_chains=config.sampling.antibody_chains),
                MergeChains(),
            ]
        )
        data_var = transform(structure_factory())
        data_variants.append(
            {
                "data": data_var,
                "name": f"{structure_id}-{config.mode}",
                "tag": config.mode,
                "residue_first": None,
                "residue_last": None,
            }
        )

    elif config.mode == "full":
        transform = Compose(
            [
                MaskAntibody(),
                MergeChains(),
            ]
        )
        data_var = transform(structure_factory())
        data_variants.append(
            {
                "data": data_var,
                "name": f"{structure_id}-Full",
                "tag": "Full",
                "residue_first": None,
                "residue_last": None,
            }
        )
    elif config.mode == "abopt":
        cdrs = sorted(
            list(set(find_cdrs(structure)).intersection(config.sampling.cdrs))
        )
        for cdr_name in cdrs:
            transform = Compose(
                [
                    MaskSingleCDR(cdr_name, augmentation=False),
                    MergeChains(),
                ]
            )
            data_var = transform(structure_factory())
            residue_first, residue_last = get_residue_first_last(data_var)
            for opt_step in config.sampling.optimize_steps:
                data_variants.append(
                    {
                        "data": data_var,
                        "name": f"{structure_id}-{cdr_name}-O{opt_step}",
                        "tag": f"{cdr_name}-O{opt_step}",
                        "cdr": cdr_name,
                        "opt_step": opt_step,
                        "residue_first": residue_first,
                        "residue_last": residue_last,
                    }
                )
    else:
        raise ValueError(f"Unknown mode: {config.mode}.")
    return data_variants


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--index", type=int, default=0)
    parser.add_argument(
        "-ck",
        "--ckpt",
        type=str,
        default="logs/seq_design__antigen-20/checkpoints/300000.pt",
    )
    parser.add_argument(
        "-c", "--config", type=str, default="configs/test/seq_design.yml"
    )
    parser.add_argument("-o", "--out_root", type=str, default="seq_design_results/")
    parser.add_argument("-r", "--relax", action="store_true", default=False)
    parser.add_argument("-n", "--num_samples", type=int, default=10)
    parser.add_argument("-t", "--tag", type=str, default="antigen-20")
    parser.add_argument("-s", "--seed", type=int, default=None)
    parser.add_argument("-d", "--device", type=str, default="cuda")
    parser.add_argument("-b", "--batch_size", type=int, default=16)
    parser.add_argument(
        "-e",
        "--eval_all",
        action="store_true",
        default=False,
        help="Call `eval_all` fn to evaluate all structures in the testset. ",
    )
    parser.add_argument(
        "--topk",
        type=int,
        default=1,
        help="topk for metric. must be less than or equal to num_samples.",
    )
    args = parser.parse_args()
    assert args.topk <= args.num_samples
    return args


def eval_all(args):
    """
    args contains arguments in parse_args() except for index, which is used to iterate over the testset.

    saved dataframe is like:

            DockQ   irms    Lrms   fnat  AAR
    H_CDR3  0.205  4.972  12.332  0.173  1.0
    retrun dict is like:
        - variant
            - metric
                - value
    """
    config, config_name = load_config(args.config)
    dataset = get_dataset(config.dataset.test)
    dname = os.path.join(
        args.out_root, config_name + "_%s" % args.tag if args.tag else ""
    )

    testset_results = {}
    results = []
    for i in range(len(dataset)):
        args.index = i
        result_dict = main(args)
        testset_results[i] = result_dict
        results.append(result_dict)
    testset_results = combine_nested_dicts(results)
    mean_result = traverse_dict(
        testset_results, list, lambda x: round(float(np.mean(x)), 3), lambda x: x
    )
    pd.DataFrame.from_dict(mean_result, orient="index").to_csv(
        os.path.join(dname, "testset_results.csv")
    )
    return mean_result


def main(args):
    """
    Args:
        args: arguments from parse_args()
    Returns:
        mean_results (dict):
            - variant tag, e.g., 'H_CDR3'
                - metric name, e.g., 'DockQ', 'AAR'
                    - metric value
    """

    # Load configs
    config, config_name = load_config(args.config)

    seed_all(args.seed if args.seed is not None else config.sampling.seed)

    # Testset
    dataset = get_dataset(config.dataset.test)
    get_structure = lambda: dataset[args.index]

    # Logging
    structure_ = get_structure()
    structure_id = structure_["id"]
    dname = os.path.join(
        args.out_root, config_name + "_%s" % args.tag if args.tag else ""
    )
    log_dir = get_new_log_dir(dname, prefix="%04d_%s" % (args.index, structure_["id"]))
    logger = get_logger("sample", log_dir)
    logger.info("Data ID: %s" % structure_["id"])
    data_native = MergeChains()(structure_)
    save_pdb(data_native, os.path.join(log_dir, "reference.pdb"))

    # Load checkpoint and model
    logger.info("Loading model config and checkpoints: %s" % (args.ckpt))
    ckpt = torch.load(args.ckpt, map_location="cpu")
    cfg_ckpt = ckpt["config"]
    model_cfg = cfg_ckpt.model
    model = get_model(model_cfg).to(args.device)
    lsd = model.load_state_dict(ckpt["model"])
    logger.info(str(lsd))

    # Make data variants
    data_variants = create_data_variants(
        config=config,
        structure_factory=get_structure,
    )

    # Save metadata
    metadata = {
        "identifier": structure_id,
        "index": args.index,
        "config": args.config,
        "items": [
            {kk: vv for kk, vv in var.items() if kk != "data"} for var in data_variants
        ],
    }
    with open(os.path.join(log_dir, "metadata.json"), "w") as f:
        json.dump(metadata, f, indent=2)

    # Start sampling
    collate_fn = PaddingCollate(eight=False)
    inference_tfm = [
        PatchAroundAnchor(
            initial_patch_size=model_cfg.initial_patch_size,
            antigen_size=model_cfg.antigen_size,
            remove_anchor=model_cfg.remove_anchor,
            crop_contiguous_antigen=model_cfg.get("crop_contiguous_antigen", False),
            contiguous_threshold=model_cfg.get("contiguous_threshold", 1e6),
        ),
    ]
    if "abopt" not in config.mode:  # Don't remove native CDR in optimization mode
        inference_tfm.append(
            RemoveNative(
                remove_structure=config.sampling.sample_structure,
                remove_sequence=config.sampling.sample_sequence,
            )
        )
    inference_tfm = Compose(inference_tfm)
    result_dict = {}
    aa_df = pd.DataFrame(columns=["Region", "native_aa", "sampled_aa", "AAR" , "PPL"])
    for variant in data_variants:
        variant_result_dict = defaultdict(list)
        os.makedirs(os.path.join(log_dir, variant["tag"]), exist_ok=True)
        logger.info(f"Start sampling for: {variant['tag']}")

        data_cropped = inference_tfm(copy.deepcopy(variant["data"]))
        data_list_repeat = [data_cropped] * args.num_samples
        loader = DataLoader(
            data_list_repeat,
            batch_size=args.batch_size,
            shuffle=False,
            collate_fn=collate_fn,
        )

        count = 0
        candidates = []
        for batch in tqdm(loader, desc=variant["name"], dynamic_ncols=True):
            torch.set_grad_enabled(False)
            model.eval()
            batch = recursive_to(batch, args.device)
            if "abopt" in config.mode:
                # Antibody optimization starting from native
                traj_batch = model.optimize(
                    batch,
                    opt_step=variant["opt_step"],
                    optimize_opt={
                        "pbar": True,
                        "sample_structure": config.sampling.sample_structure,
                        "sample_sequence": config.sampling.sample_sequence,
                    },
                )
            else:
                # De novo design
                traj_batch = model.sample(
                    batch,
                    sample_opt={
                        "pbar": True,
                        "sample_structure": config.sampling.sample_structure,
                        "sample_sequence": config.sampling.sample_sequence,
                        "contig":"", # do not specify contig regions on test set evaluation.
                    },
                )

            aa_new = traj_batch[0][2]  # 0: Last sampling step. 2: Amino acid.
            prmsd = traj_batch[0][3]  # 3: prmsd, [B, N]
            perplexity = traj_batch[0][4]  # 4: perplexity, [B,]
            pos_atom_new, mask_atom_new = reconstruct_backbone_partially(
                pos_ctx=batch["pos_heavyatom"],
                R_new=so3vec_to_rotation(traj_batch[0][0]),
                t_new=traj_batch[0][1],
                aa=aa_new,
                chain_nb=batch["chain_nb"],
                res_nb=batch["res_nb"],
                mask_atoms=batch["mask_heavyatom"],
                mask_recons=batch["generate_flag"],
            )
            aa_new = aa_new.cpu()
            pos_atom_new = pos_atom_new.cpu()
            mask_atom_new = mask_atom_new.cpu()

            # save ref docking pose
            data_tmpl = variant["data"]
            antigen_mask = data_tmpl["fragment_type"] == constants.Fragment.Antigen
            patch_mask = torch.zeros_like(antigen_mask).bool()
            patch_mask[data_cropped["patch_idx"]] = True
            antigen_and_patch_mask = antigen_mask | patch_mask
            native_patch = _mask_select_data(data_tmpl, antigen_and_patch_mask)
            ref_path = os.path.join(log_dir, variant["tag"], "REF1.pdb")
            save_pdb(native_patch, path=ref_path)

            for i in range(aa_new.size(0)):
                data_tmpl = variant["data"]
                generate_flag_i = batch["generate_flag"][i]
                prmsd_i = round(
                    float(prmsd[i].mean()), 6
                )  # round(float(prmsd[i][generate_flag_i].mean()),3)
                perplexity_i = round(float(perplexity[i]), 6)

                pred_aa = aa_new[i][generate_flag_i]
                aa = apply_patch_to_tensor(
                    data_tmpl["aa"], aa_new[i], data_cropped["patch_idx"]
                )
                mask_ha = apply_patch_to_tensor(
                    data_tmpl["mask_heavyatom"],
                    mask_atom_new[i],
                    data_cropped["patch_idx"],
                )
                pos_ha = apply_patch_to_tensor(
                    data_tmpl["pos_heavyatom"],
                    pos_atom_new[i] + batch["origin"][i].view(1, 1, 3).cpu(),
                    data_cropped["patch_idx"],
                )
                native_aa = data_tmpl["aa"][data_tmpl["generate_flag"]]
                generated_structure = pos_ha[data_tmpl["generate_flag"]]
                candidates.append(generated_structure)
                aar = float((native_aa == pred_aa).sum().float() / native_aa.size(0))
                native_aa_codes = "".join(
                    constants.aaidx2symbol[int(i)] for i in native_aa
                )
                pred_aa_codes = "".join(constants.aaidx2symbol[int(i)] for i in pred_aa)
                aa_df = pd.concat(
                    [
                        aa_df,
                        pd.DataFrame(
                            {
                                "Region": [variant["tag"]],
                                "native_aa": [native_aa_codes],
                                "sampled_aa": [pred_aa_codes],
                                "AAR": [aar],
                                "PPL": [perplexity_i],
                            }
                        ),
                    ]
                )
                # save cdr + antigen complex.
                antigen_mask = data_tmpl["fragment_type"] == constants.Fragment.Antigen
                patch_mask = torch.zeros_like(antigen_mask).bool()
                patch_mask[data_cropped["patch_idx"]] = True
                antigen_and_patch_mask = antigen_mask | patch_mask
                data = copy.deepcopy(data_tmpl)
                data["aa"] = aa
                data["mask_heavyatom"] = mask_ha
                data["pos_heavyatom"] = pos_ha
                data_patch = _mask_select_data(data, antigen_and_patch_mask)
                save_path = os.path.join(log_dir, variant["tag"], "%04d.pdb" % (count,))
                save_pdb(data_patch, path=save_path)
                del data, data_patch
                # save_pdb({
                #     'chain_nb': data_tmpl['chain_nb'],
                #     'chain_id': data_tmpl['chain_id'],
                #     'resseq': data_tmpl['resseq'],
                #     'icode': data_tmpl['icode'],
                #     # Generated
                #     'aa': aa,
                #     'mask_heavyatom': mask_ha,
                #     'pos_heavyatom': pos_ha,
                # }, path=save_path)

                patch_path = os.path.join(
                    log_dir, variant["tag"], "%04d_patch.pdb" % (count,)
                )
                save_pdb(
                    {
                        "chain_nb": data_cropped["chain_nb"],
                        "chain_id": data_cropped["chain_id"],
                        "resseq": data_cropped["resseq"],
                        "icode": data_cropped["icode"],
                        # Generated
                        "aa": aa_new[i],
                        "mask_heavyatom": mask_atom_new[i],
                        "pos_heavyatom": pos_atom_new[i]
                        + batch["origin"][i].view(1, 1, 3).cpu(),
                    },
                    path=patch_path,
                )
                count += 1

                score_dict = calc_DockQ(save_path, ref_path, use_CA_only=True)
                score_dict = {
                    k: round(v, 3)
                    for k, v in score_dict.items()
                    if k in ["DockQ", "irms", "Lrms", "fnat"]
                }
                score_dict.update({"AAR": aar, "prmsd": prmsd_i})
                [variant_result_dict[key].append(score_dict[key]) for key in score_dict]
            logger.info("Finished.\n")
        candidates = torch.stack(candidates, dim=0)  # (n_candidates, n_res, n_atoms, 3)
        candidates = candidates[
            ...,
            (
                constants.BBHeavyAtom.N,
                constants.BBHeavyAtom.CA,
                constants.BBHeavyAtom.CB,
            ),
            :,
        ]
        candidates = candidates.reshape(count, -1, 3)
        topk_idxs = rank_commoness(candidates, args.topk)
        top_result = {
            f"{k}_top{args.topk}": [v[i] for i in topk_idxs]
            for k, v in variant_result_dict.items()
        }
        variant_result_dict.update(top_result)
        result_dict[variant["tag"]] = variant_result_dict
    results = traverse_dict(
        result_dict, list, lambda x: round(float(np.mean(x)), 3), lambda x: f"{x}_mean"
    )
    # mean_results = traverse_dict(result_dict, list, lambda x: round(float(np.mean(x)),3), lambda x: x)
    topk_results = extract_dict(
        result_dict, list, lambda k: k.endswith(f"_top{args.topk}")
    )
    df = pd.DataFrame.from_dict(results, orient="index")
    df.to_csv(os.path.join(log_dir, "results.csv"))
    aa_df.to_csv(os.path.join(log_dir, "aa.csv"), index=False)
    if args.relax:
        run_relax(log_dir, pipeline_name="openmm_pyrosetta")
        run_energy(log_dir)
        openmm_path = save_path.replace(".pdb", "_openmm.pdb")
        rosetta_path = save_path.replace(".pdb", "_rosetta.pdb")
    return topk_results  # mean_results


def traverse_dict(d, leaf_type, leaf_fn, leaf_key_fn, leaf_key_vali_fn=lambda x: True):
    """
    traverse a dict and apply a function to all leaves if they are of a certain type
    return a new dict with the same structure as the input dict, but with the leaves replaced by the result of the function.
    :param d: dict to traverse
    :param leaf_type: type of the leaves to apply the function to
    :param leaf_fn: function to apply to the leaves
    :return: new dict with the same structure as the input dict, but with the leaves replaced by the result of the function.
    """
    ret_dict = copy.deepcopy(d)
    for k, v in d.items():
        if isinstance(v, dict):
            ret_dict[k] = traverse_dict(v, leaf_type, leaf_fn, leaf_key_fn)
        elif isinstance(v, leaf_type) and leaf_key_vali_fn(k):
            ret_dict[leaf_key_fn(k)] = leaf_fn(v)
    return ret_dict


def extract_dict(d, leaf_type, leaf_key_vali_fn=lambda x: True):
    """
    extract a dict with only leaves of a certain type
    return a new dict with the same structure as the input dict, but with the leaves replaced by the result of the function.
    :param d: dict to traverse
    :param leaf_type: type of the leaves to apply the function to
    :param leaf_fn: function to apply to the leaves
    :return: new dict with the same structure as the input dict, but with the leaves replaced by the result of the function.
    """
    ret_dict = copy.deepcopy(d)
    for k, v in d.items():
        if isinstance(v, dict):
            ret_dict[k] = extract_dict(v, leaf_type, leaf_key_vali_fn)
        elif isinstance(v, leaf_type) and leaf_key_vali_fn(k):
            ret_dict[k] = v
        else:
            del ret_dict[k]
    return ret_dict


def combine_nested_dicts(dicts):
    """
    Combine a list of nested dictionaries into a single nested dictionary.

    Args:
        dicts (list): List of nested dictionaries with similar keys.

    Returns:
        dict: A nested dictionary with the same structure as the input dictionaries.
    """
    if len(dicts) == 1:
        return dicts[0]
    else:
        combined_dict = {}
        for key in dicts[0].keys():
            if isinstance(dicts[0][key], dict):
                combined_dict[key] = combine_nested_dicts([d[key] for d in dicts])
            else:
                combined_dict[key] = [d[key] for d in dicts]
        return combined_dict


def calc_per_rmsd(structures):
    B, N, _ = structures.shape
    structures = structures.unsqueeze(1).repeat(1, B, 1, 1)  # (B, B, N, 3)
    structures_ = structures.permute(1, 0, 2, 3)  # (B, B, N, 3)
    rmsd = torch.sqrt(
        (((structures - structures_) ** 2).sum(dim=-1)).mean(-1)
    )  # (B, B)
    return rmsd


def calc_avg_rmsd(structures):
    B, N, _ = structures.shape
    rmsd = calc_per_rmsd(structures)
    avg_rmsd = rmsd.sum() / (B * (B - 1))
    return avg_rmsd


def rank_commoness(structures, k):
    """
    Find out the most common structure within a list of structures.
    A structure is common if it is similar to all other structures.
    The similarity is measured by RMSD. The lower the RMSD, the more similar the structures are.
    Args:
        structures(torch.Tensor): (B, N, 3). B is num of structures, N is the structure length.
    Returns:
        rank (torch.Tensor): (B, ), the similarity rank of each structure.
    """
    B, N, _ = structures.shape
    rmsd = calc_per_rmsd(structures)  # (B, B)
    similarity, rank = torch.topk(
        rmsd.sum(dim=-1) / (B - 1), k=k, largest=False
    )  # (B, )

    return rank

AbDock/src/tools/runner/design_for_testset.py
pengzhangzhi/ab_opt
import math
import torch
from torch.utils.data._utils.collate import default_collate


DEFAULT_PAD_VALUES = {
    'aa': 21, 
    'chain_id': ' ', 
    'icode': ' ',
}

DEFAULT_NO_PADDING = {
    'origin',
}

class PaddingCollate(object):

    def __init__(self, length_ref_key='aa', pad_values=DEFAULT_PAD_VALUES, no_padding=DEFAULT_NO_PADDING, eight=True):
        super().__init__()
        self.length_ref_key = length_ref_key
        self.pad_values = pad_values
        self.no_padding = no_padding
        self.eight = eight

    @staticmethod
    def _pad_last(x, n, value=0):
        if isinstance(x, torch.Tensor):
            assert x.size(0) <= n
            if x.size(0) == n:
                return x
            pad_size = [n - x.size(0)] + list(x.shape[1:])
            pad = torch.full(pad_size, fill_value=value).to(x)
            return torch.cat([x, pad], dim=0)
        elif isinstance(x, list):
            pad = [value] * (n - len(x))
            return x + pad
        else:
            return x

    @staticmethod
    def _get_pad_mask(l, n):
        return torch.cat([
            torch.ones([l], dtype=torch.bool),
            torch.zeros([n-l], dtype=torch.bool)
        ], dim=0)

    @staticmethod
    def _get_common_keys(list_of_dict):
        keys = set(list_of_dict[0].keys())
        for d in list_of_dict[1:]:
            keys = keys.intersection(d.keys())
        return keys


    def _get_pad_value(self, key):
        if key not in self.pad_values:
            return 0
        return self.pad_values[key]

    def __call__(self, data_list):
        max_length = max([data[self.length_ref_key].size(0) for data in data_list])
        keys = self._get_common_keys(data_list)
        
        if self.eight:
            max_length = math.ceil(max_length / 8) * 8
        data_list_padded = []
        for data in data_list:
            data_padded = {
                k: self._pad_last(v, max_length, value=self._get_pad_value(k)) if k not in self.no_padding else v
                for k, v in data.items()
                if k in keys
            }
            data_padded['mask'] = self._get_pad_mask(data[self.length_ref_key].size(0), max_length)
            data_list_padded.append(data_padded)
        return default_collate(data_list_padded)


def apply_patch_to_tensor(x_full, x_patch, patch_idx):
    """
    Args:
        x_full:  (N, ...)
        x_patch: (M, ...)
        patch_idx:  (M, )
    Returns:
        (N, ...)
    """
    x_full = x_full.clone()
    x_full[patch_idx] = x_patch
    return x_full

AbDock/src/utils/data.py
pengzhangzhi/ab_opt
import torch
from .protein import constants


def find_cdrs(structure):
    cdrs = []
    if structure['heavy'] is not None:
        flag = structure['heavy']['cdr_flag']
        if int(constants.CDR.H1) in flag:
            cdrs.append('H_CDR1')
        if int(constants.CDR.H2) in flag:
            cdrs.append('H_CDR2')
        if int(constants.CDR.H3) in flag:
            cdrs.append('H_CDR3')

    if structure['light'] is not None:
        flag = structure['light']['cdr_flag']
        if int(constants.CDR.L1) in flag:
            cdrs.append('L_CDR1')
        if int(constants.CDR.L2) in flag:
            cdrs.append('L_CDR2')
        if int(constants.CDR.L3) in flag:
            cdrs.append('L_CDR3')
    
    return cdrs


def get_residue_first_last(data):
    loop_flag = data['generate_flag']
    loop_idx = torch.arange(loop_flag.size(0))[loop_flag]
    idx_first, idx_last = loop_idx.min().item(), loop_idx.max().item()
    residue_first = (data['chain_id'][idx_first], data['resseq'][idx_first].item(), data['icode'][idx_first])
    residue_last = (data['chain_id'][idx_last], data['resseq'][idx_last].item(), data['icode'][idx_last])
    return residue_first, residue_last


class RemoveNative(object):

    def __init__(self, remove_structure, remove_sequence):
        super().__init__()
        self.remove_structure = remove_structure
        self.remove_sequence = remove_sequence

    def __call__(self, data):
        generate_flag = data['generate_flag'].clone()
        if self.remove_sequence:
            data['aa'] = torch.where(
                generate_flag, 
                torch.full_like(data['aa'], fill_value=int(constants.AA.UNK)),    # Is loop
                data['aa']
            )

        if self.remove_structure:
            data['pos_heavyatom'] = torch.where(
                generate_flag[:, None, None].expand(data['pos_heavyatom'].shape),
                torch.randn_like(data['pos_heavyatom']) * 10,
                data['pos_heavyatom']
            )

        return data
AbDock/src/utils/inference.py
pengzhangzhi/ab_opt
import os
import time
import random
import logging
from typing import OrderedDict
import torch
import torch.linalg
import numpy as np
import yaml
from easydict import EasyDict
from glob import glob
import dynamic_yaml
import pdbfixer
import openmm
ENERGY = openmm.unit.kilocalories_per_mole
LENGTH = openmm.unit.angstroms

class BlackHole(object):
    def __setattr__(self, name, value):
        pass

    def __call__(self, *args, **kwargs):
        return self

    def __getattr__(self, name):
        return self


class Counter(object):
    def __init__(self, start=0):
        super().__init__()
        self.now = start

    def step(self, delta=1):
        prev = self.now
        self.now += delta
        return prev


def get_logger(name, log_dir=None):
    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter('[%(asctime)s::%(name)s::%(levelname)s] %(message)s')

    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(logging.DEBUG)
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)

    if log_dir is not None:
        file_handler = logging.FileHandler(os.path.join(log_dir, 'log.txt'))
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

    return logger


def get_new_log_dir(root='./logs', prefix='', tag=''):
    fn = "" # time.strftime('%Y_%m_%d__%H_%M_%S', time.localtime())
    if prefix != '':
        fn = prefix + '_' + fn
    if tag != '':
        fn = fn + '_' + tag
    log_dir = os.path.join(root, fn)
    os.makedirs(log_dir, exist_ok=True)
    return log_dir


def seed_all(seed):
    torch.backends.cudnn.deterministic = True
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)


def inf_iterator(iterable):
    iterator = iterable.__iter__()
    while True:
        try:
            yield iterator.__next__()
        except StopIteration:
            iterator = iterable.__iter__()


def log_hyperparams(writer, args):
    from torch.utils.tensorboard.summary import hparams
    vars_args = {k: v if isinstance(v, str) else repr(v) for k, v in vars(args).items()}
    exp, ssi, sei = hparams(vars_args, {})
    writer.file_writer.add_summary(exp)
    writer.file_writer.add_summary(ssi)
    writer.file_writer.add_summary(sei)


def int_tuple(argstr):
    return tuple(map(int, argstr.split(',')))


def str_tuple(argstr):
    return tuple(argstr.split(','))


def get_checkpoint_path(folder, it=None):
    if it is not None:
        return os.path.join(folder, '%d.pt' % it), it
    all_iters = list(map(lambda x: int(os.path.basename(x[:-3])), glob(os.path.join(folder, '*.pt'))))
    all_iters.sort()
    return os.path.join(folder, '%d.pt' % all_iters[-1]), all_iters[-1]


def load_config(config_path):
    with open(config_path, 'r') as f:
        config = EasyDict(dynamic_yaml.load(f))
    config_name = os.path.basename(config_path)[:os.path.basename(config_path).rfind('.')]
    return config, config_name


def extract_weights(weights: OrderedDict, prefix):
    extracted = OrderedDict()
    for k, v in weights.items():
        if k.startswith(prefix):
            extracted.update({
                k[len(prefix):]: v
            })
    return extracted


def current_milli_time():
    return round(time.time() * 1000)



def openmm_relax(pdb_file, stiffness=10., tolerance=2.39, use_gpu=False):
    fixer = pdbfixer.PDBFixer(pdb_file)
    fixer.findMissingResidues()
    fixer.findMissingAtoms()
    fixer.addMissingAtoms()
    fixer.addMissingHydrogens()

    force_field = openmm.app.ForceField("amber14/protein.ff14SB.xml")
    modeller = openmm.app.Modeller(fixer.topology, fixer.positions)
    modeller.addHydrogens(force_field)
    system = force_field.createSystem(modeller.topology)

    if stiffness > 0:
        stiffness = stiffness * ENERGY / (LENGTH**2)
        force = openmm.CustomExternalForce("0.5 * k * ((x-x0)^2 + (y-y0)^2 + (z-z0)^2)")
        force.addGlobalParameter("k", stiffness)
        for p in ["x0", "y0", "z0"]:
            force.addPerParticleParameter(p)
        for residue in modeller.topology.residues():
            for atom in residue.atoms():
                if atom.name in ["N", "CA", "C", "CB"]:
                    force.addParticle(
                            atom.index, modeller.positions[atom.index]
                    )
        system.addForce(force)

    tolerance = tolerance * ENERGY
    integrator = openmm.LangevinIntegrator(0, 0.01, 1.0)
    platform = openmm.Platform.getPlatformByName("CUDA" if use_gpu else "CPU")

    simulation = openmm.app.Simulation(modeller.topology, system, integrator, platform)
    simulation.context.setPositions(modeller.positions)
    simulation.minimizeEnergy(tolerance)
    state = simulation.context.getState(getEnergy=True)
    energy = state.getKineticEnergy() + state.getPotentialEnergy()

    with open(pdb_file, "w") as f:
        openmm.app.PDBFile.writeFile(
            simulation.topology,
            simulation.context.getState(getPositions=True).getPositions(),
            f,
            keepIds=True
        )
    return energy


AbDock/src/utils/misc.py
pengzhangzhi/ab_opt
import torch
import enum

class CDR(enum.IntEnum):
    H1 = 1
    H2 = 2
    H3 = 3
    L1 = 4
    L2 = 5
    L3 = 6


class ChothiaCDRRange:
    H1 = (26, 32)
    H2 = (52, 56)
    H3 = (95, 102)
    
    L1 = (24, 34)
    L2 = (50, 56)
    L3 = (89, 97)

    @classmethod
    def to_cdr(cls, chain_type, resseq):
        assert chain_type in ('H', 'L')
        if chain_type == 'H':
            if cls.H1[0] <= resseq <= cls.H1[1]:
                return CDR.H1
            elif cls.H2[0] <= resseq <= cls.H2[1]:
                return CDR.H2
            elif cls.H3[0] <= resseq <= cls.H3[1]:
                return CDR.H3
        elif chain_type == 'L':
            if cls.L1[0] <= resseq <= cls.L1[1]:     # Chothia VH-CDR1
                return CDR.L1
            elif cls.L2[0] <= resseq <= cls.L2[1]:
                return CDR.L2
            elif cls.L3[0] <= resseq <= cls.L3[1]:
                return CDR.L3


class Fragment(enum.IntEnum):
    Heavy = 1
    Light = 2
    Antigen = 3

##
# Residue identities
"""
This is part of the OpenMM molecular simulation toolkit originating from
Simbios, the NIH National Center for Physics-Based Simulation of
Biological Structures at Stanford, funded under the NIH Roadmap for
Medical Research, grant U54 GM072970. See https://simtk.org.

Portions copyright (c) 2013 Stanford University and the Authors.
Authors: Peter Eastman
Contributors:

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
non_standard_residue_substitutions = {
    '2AS':'ASP', '3AH':'HIS', '5HP':'GLU', 'ACL':'ARG', 'AGM':'ARG', 'AIB':'ALA', 'ALM':'ALA', 'ALO':'THR', 'ALY':'LYS', 'ARM':'ARG',
    'ASA':'ASP', 'ASB':'ASP', 'ASK':'ASP', 'ASL':'ASP', 'ASQ':'ASP', 'AYA':'ALA', 'BCS':'CYS', 'BHD':'ASP', 'BMT':'THR', 'BNN':'ALA',
    'BUC':'CYS', 'BUG':'LEU', 'C5C':'CYS', 'C6C':'CYS', 'CAS':'CYS', 'CCS':'CYS', 'CEA':'CYS', 'CGU':'GLU', 'CHG':'ALA', 'CLE':'LEU', 'CME':'CYS',
    'CSD':'ALA', 'CSO':'CYS', 'CSP':'CYS', 'CSS':'CYS', 'CSW':'CYS', 'CSX':'CYS', 'CXM':'MET', 'CY1':'CYS', 'CY3':'CYS', 'CYG':'CYS',
    'CYM':'CYS', 'CYQ':'CYS', 'DAH':'PHE', 'DAL':'ALA', 'DAR':'ARG', 'DAS':'ASP', 'DCY':'CYS', 'DGL':'GLU', 'DGN':'GLN', 'DHA':'ALA',
    'DHI':'HIS', 'DIL':'ILE', 'DIV':'VAL', 'DLE':'LEU', 'DLY':'LYS', 'DNP':'ALA', 'DPN':'PHE', 'DPR':'PRO', 'DSN':'SER', 'DSP':'ASP',
    'DTH':'THR', 'DTR':'TRP', 'DTY':'TYR', 'DVA':'VAL', 'EFC':'CYS', 'FLA':'ALA', 'FME':'MET', 'GGL':'GLU', 'GL3':'GLY', 'GLZ':'GLY',
    'GMA':'GLU', 'GSC':'GLY', 'HAC':'ALA', 'HAR':'ARG', 'HIC':'HIS', 'HIP':'HIS', 'HMR':'ARG', 'HPQ':'PHE', 'HTR':'TRP', 'HYP':'PRO',
    'IAS':'ASP', 'IIL':'ILE', 'IYR':'TYR', 'KCX':'LYS', 'LLP':'LYS', 'LLY':'LYS', 'LTR':'TRP', 'LYM':'LYS', 'LYZ':'LYS', 'MAA':'ALA', 'MEN':'ASN',
    'MHS':'HIS', 'MIS':'SER', 'MLE':'LEU', 'MPQ':'GLY', 'MSA':'GLY', 'MSE':'MET', 'MVA':'VAL', 'NEM':'HIS', 'NEP':'HIS', 'NLE':'LEU',
    'NLN':'LEU', 'NLP':'LEU', 'NMC':'GLY', 'OAS':'SER', 'OCS':'CYS', 'OMT':'MET', 'PAQ':'TYR', 'PCA':'GLU', 'PEC':'CYS', 'PHI':'PHE',
    'PHL':'PHE', 'PR3':'CYS', 'PRR':'ALA', 'PTR':'TYR', 'PYX':'CYS', 'SAC':'SER', 'SAR':'GLY', 'SCH':'CYS', 'SCS':'CYS', 'SCY':'CYS',
    'SEL':'SER', 'SEP':'SER', 'SET':'SER', 'SHC':'CYS', 'SHR':'LYS', 'SMC':'CYS', 'SOC':'CYS', 'STY':'TYR', 'SVA':'SER', 'TIH':'ALA',
    'TPL':'TRP', 'TPO':'THR', 'TPQ':'ALA', 'TRG':'LYS', 'TRO':'TRP', 'TYB':'TYR', 'TYI':'TYR', 'TYQ':'TYR', 'TYS':'TYR', 'TYY':'TYR'
}


ressymb_to_resindex = {
    'A': 0, 'C': 1, 'D': 2, 'E': 3, 'F': 4,
    'G': 5, 'H': 6, 'I': 7, 'K': 8, 'L': 9,
    'M': 10, 'N': 11, 'P': 12, 'Q': 13, 'R': 14,
    'S': 15, 'T': 16, 'V': 17, 'W': 18, 'Y': 19,
    'X': 20,
}
aaidx2symbol = {v: k for k, v in ressymb_to_resindex.items()}

class AA(enum.IntEnum):
    ALA = 0; CYS = 1; ASP = 2; GLU = 3; PHE = 4
    GLY = 5; HIS = 6; ILE = 7; LYS = 8; LEU = 9
    MET = 10; ASN = 11; PRO = 12; GLN = 13; ARG = 14
    SER = 15; THR = 16; VAL = 17; TRP = 18; TYR = 19
    UNK = 20

    @classmethod
    def _missing_(cls, value):
        if isinstance(value, str) and len(value) == 3:      # three representation
            if value in non_standard_residue_substitutions:
                value = non_standard_residue_substitutions[value]
            if value in cls._member_names_:
                return getattr(cls, value)
        elif isinstance(value, str) and len(value) == 1:    # one representation
            if value in ressymb_to_resindex:
                return cls(ressymb_to_resindex[value])

        return super()._missing_(value)

    def __str__(self):
        return self.name

    @classmethod
    def is_aa(cls, value):
        return (value in ressymb_to_resindex) or \
            (value in non_standard_residue_substitutions) or \
            (value in cls._member_names_) or \
            (value in cls._member_map_.values())


num_aa_types = len(AA)

##
# Atom identities

class BBHeavyAtom(enum.IntEnum):
    N = 0; CA = 1; C = 2; O = 3; CB = 4; OXT=14;

NUM_BB_ATOMS = 5
max_num_heavyatoms = 15

# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
restype_to_heavyatom_names = {
    AA.ALA: ['N', 'CA', 'C', 'O', 'CB', '',    '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.ARG: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'NE',  'CZ',  'NH1', 'NH2', '',    '',    '', 'OXT'],
    AA.ASN: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'ND2', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.ASP: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'OD2', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.CYS: ['N', 'CA', 'C', 'O', 'CB', 'SG',  '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.GLN: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'NE2', '',    '',    '',    '',    '', 'OXT'],
    AA.GLU: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'OE2', '',    '',    '',    '',    '', 'OXT'],
    AA.GLY: ['N', 'CA', 'C', 'O', '',   '',    '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.HIS: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'ND1', 'CD2', 'CE1', 'NE2', '',    '',    '',    '', 'OXT'],
    AA.ILE: ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', 'CD1', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.LEU: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', '',    '',    '',    '',    '',    '', 'OXT'],
    AA.LYS: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'CE',  'NZ',  '',    '',    '',    '',    '', 'OXT'],
    AA.MET: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'SD',  'CE',  '',    '',    '',    '',    '',    '', 'OXT'],
    AA.PHE: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  '',    '',    '', 'OXT'],
    AA.PRO: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.SER: ['N', 'CA', 'C', 'O', 'CB', 'OG',  '',    '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.THR: ['N', 'CA', 'C', 'O', 'CB', 'OG1', 'CG2', '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.TRP: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'NE1', 'CE2', 'CE3', 'CZ2', 'CZ3', 'CH2', 'OXT'],
    AA.TYR: ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  'OH',  '',    '', 'OXT'],
    AA.VAL: ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '',    '',    '',    '',    '',    '',    '', 'OXT'],
    AA.UNK: ['',  '',   '',  '',  '',   '',    '',    '',    '',    '',    '',    '',    '',    '',    ''],
}
for names in restype_to_heavyatom_names.values(): assert len(names) == max_num_heavyatoms


backbone_atom_coordinates = {
    AA.ALA: [
        (-0.525, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, -0.0, -0.0),  # C
    ],
    AA.ARG: [
        (-0.524, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, -0.0),  # C
    ],
    AA.ASN: [
        (-0.536, 1.357, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, -0.0, -0.0),  # C
    ],
    AA.ASP: [
        (-0.525, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, 0.0, -0.0),  # C
    ],
    AA.CYS: [
        (-0.522, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.524, 0.0, 0.0),  # C
    ],
    AA.GLN: [
        (-0.526, 1.361, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, 0.0, 0.0),  # C
    ],
    AA.GLU: [
        (-0.528, 1.361, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, -0.0, -0.0),  # C
    ],
    AA.GLY: [
        (-0.572, 1.337, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.517, -0.0, -0.0),  # C
    ],
    AA.HIS: [
        (-0.527, 1.36, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, 0.0, 0.0),  # C
    ],
    AA.ILE: [
        (-0.493, 1.373, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, -0.0, -0.0),  # C
    ],
    AA.LEU: [
        (-0.52, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, -0.0),  # C
    ],
    AA.LYS: [
        (-0.526, 1.362, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, 0.0, 0.0),  # C
    ],
    AA.MET: [
        (-0.521, 1.364, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, 0.0, 0.0),  # C
    ],
    AA.PHE: [
        (-0.518, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.524, 0.0, -0.0),  # C
    ],
    AA.PRO: [
        (-0.566, 1.351, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, -0.0, 0.0),  # C
    ],
    AA.SER: [
        (-0.529, 1.36, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, -0.0),  # C
    ],
    AA.THR: [
        (-0.517, 1.364, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.526, 0.0, -0.0),  # C
    ],
    AA.TRP: [
        (-0.521, 1.363, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.525, -0.0, 0.0),  # C
    ],
    AA.TYR: [
        (-0.522, 1.362, 0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.524, -0.0, -0.0),  # C
    ],
    AA.VAL: [
        (-0.494, 1.373, -0.0),  # N
        (0.0, 0.0, 0.0),  # CA
        (1.527, -0.0, -0.0),  # C
    ],
}

bb_oxygen_coordinate = {
    AA.ALA: (2.153, -1.062, 0.0),
    AA.ARG: (2.151, -1.062, 0.0),
    AA.ASN: (2.151, -1.062, 0.0),
    AA.ASP: (2.153, -1.062, 0.0),
    AA.CYS: (2.149, -1.062, 0.0),
    AA.GLN: (2.152, -1.062, 0.0),
    AA.GLU: (2.152, -1.062, 0.0),
    AA.GLY: (2.143, -1.062, 0.0),
    AA.HIS: (2.15, -1.063, 0.0),
    AA.ILE: (2.154, -1.062, 0.0),
    AA.LEU: (2.15, -1.063, 0.0),
    AA.LYS: (2.152, -1.062, 0.0),
    AA.MET: (2.15, -1.062, 0.0),
    AA.PHE: (2.15, -1.062, 0.0),
    AA.PRO: (2.148, -1.066, 0.0),
    AA.SER: (2.151, -1.062, 0.0),
    AA.THR: (2.152, -1.062, 0.0),
    AA.TRP: (2.152, -1.062, 0.0),
    AA.TYR: (2.151, -1.062, 0.0),
    AA.VAL: (2.154, -1.062, 0.0),
}

backbone_atom_coordinates_tensor = torch.zeros([21, 3, 3])
bb_oxygen_coordinate_tensor = torch.zeros([21, 3])

def make_coordinate_tensors():
    for restype, atom_coords in backbone_atom_coordinates.items():
        for atom_id, atom_coord in enumerate(atom_coords):
            backbone_atom_coordinates_tensor[restype][atom_id] = torch.FloatTensor(atom_coord)
    
    for restype, bb_oxy_coord in bb_oxygen_coordinate.items():
        bb_oxygen_coordinate_tensor[restype] = torch.FloatTensor(bb_oxy_coord)
make_coordinate_tensors()

AbDock/src/utils/protein/constants.py
pengzhangzhi/ab_opt
import torch
from Bio.PDB import Selection
from Bio.PDB.Residue import Residue
from easydict import EasyDict

from .constants import (
    AA, max_num_heavyatoms,
    restype_to_heavyatom_names, 
    BBHeavyAtom
)


class ParsingException(Exception):
    pass


def _get_residue_heavyatom_info(res: Residue):
    pos_heavyatom = torch.zeros([max_num_heavyatoms, 3], dtype=torch.float)
    mask_heavyatom = torch.zeros([max_num_heavyatoms, ], dtype=torch.bool)
    restype = AA(res.get_resname())
    for idx, atom_name in enumerate(restype_to_heavyatom_names[restype]):
        if atom_name == '': continue
        if atom_name in res:
            pos_heavyatom[idx] = torch.tensor(res[atom_name].get_coord().tolist(), dtype=pos_heavyatom.dtype)
            mask_heavyatom[idx] = True
    return pos_heavyatom, mask_heavyatom


def parse_biopython_structure(entity, unknown_threshold=1.0, max_resseq=None):
    chains = Selection.unfold_entities(entity, 'C')
    chains.sort(key=lambda c: c.get_id())
    data = EasyDict({
        'chain_id': [],
        'resseq': [], 'icode': [], 'res_nb': [],
        'aa': [],
        'pos_heavyatom': [], 'mask_heavyatom': [],
    })
    tensor_types = {
        'resseq': torch.LongTensor,
        'res_nb': torch.LongTensor,
        'aa': torch.LongTensor,
        'pos_heavyatom': torch.stack,
        'mask_heavyatom': torch.stack,
    }

    count_aa, count_unk = 0, 0

    for i, chain in enumerate(chains):
        seq_this = 0   # Renumbering residues
        residues = Selection.unfold_entities(chain, 'R')
        residues.sort(key=lambda res: (res.get_id()[1], res.get_id()[2]))   # Sort residues by resseq-icode
        for _, res in enumerate(residues):
            resseq_this = int(res.get_id()[1])
            if max_resseq is not None and resseq_this > max_resseq:
                continue

            resname = res.get_resname()
            if not AA.is_aa(resname): continue
            if not (res.has_id('CA') and res.has_id('C') and res.has_id('N')): continue
            restype = AA(resname)
            count_aa += 1
            if restype == AA.UNK: 
                count_unk += 1
                continue

            # Chain info
            data.chain_id.append(chain.get_id())

            # Residue types
            data.aa.append(restype) # Will be automatically cast to torch.long

            # Heavy atoms
            pos_heavyatom, mask_heavyatom = _get_residue_heavyatom_info(res)
            data.pos_heavyatom.append(pos_heavyatom)
            data.mask_heavyatom.append(mask_heavyatom)

            # Sequential number
            resseq_this = int(res.get_id()[1])
            icode_this = res.get_id()[2]
            if seq_this == 0:
                seq_this = 1
            else:
                d_CA_CA = torch.linalg.norm(data.pos_heavyatom[-2][BBHeavyAtom.CA] - data.pos_heavyatom[-1][BBHeavyAtom.CA], ord=2).item()
                if d_CA_CA <= 4.0:
                    seq_this += 1
                else:
                    d_resseq = resseq_this - data.resseq[-1]
                    seq_this += max(2, d_resseq)

            data.resseq.append(resseq_this)
            data.icode.append(icode_this)
            data.res_nb.append(seq_this)

    if len(data.aa) == 0:
        raise ParsingException('No parsed residues.')

    if (count_unk / count_aa) >= unknown_threshold:
        raise ParsingException(
            f'Too many unknown residues, threshold {unknown_threshold:.2f}.'
        )

    seq_map = {}
    for i, (chain_id, resseq, icode) in enumerate(zip(data.chain_id, data.resseq, data.icode)):
        seq_map[(chain_id, resseq, icode)] = i

    for key, convert_fn in tensor_types.items():
        data[key] = convert_fn(data[key])

    return data, seq_map

AbDock/src/utils/protein/parsers.py
pengzhangzhi/ab_opt
import torch
import warnings
from Bio import BiopythonWarning
from Bio.PDB import PDBIO
from Bio.PDB.StructureBuilder import StructureBuilder

from .constants import AA, restype_to_heavyatom_names


def save_pdb(data, path=None):
    """
    Args:
        data:   A dict that contains: `chain_nb`, `chain_id`, `aa`, `resseq`, `icode`,
                `pos_heavyatom`, `mask_heavyatom`.
    """

    def _mask_select(v, mask):
        if isinstance(v, str):
            return ''.join([s for i, s in enumerate(v) if mask[i]])
        elif isinstance(v, list):
            return [s for i, s in enumerate(v) if mask[i]]
        elif isinstance(v, torch.Tensor):
            return v[mask]
        else:
            return v

    def _build_chain(builder, aa_ch, pos_heavyatom_ch, mask_heavyatom_ch, chain_id_ch, resseq_ch, icode_ch):
        builder.init_chain(chain_id_ch[0])
        builder.init_seg('    ')

        for aa_res, pos_allatom_res, mask_allatom_res, resseq_res, icode_res in \
            zip(aa_ch, pos_heavyatom_ch, mask_heavyatom_ch, resseq_ch, icode_ch):
            if not AA.is_aa(aa_res.item()): 
                print('[Warning] Unknown amino acid type at %d%s: %r' % (resseq_res.item(), icode_res, aa_res.item()))
                continue
            restype = AA(aa_res.item())
            builder.init_residue(
                resname = str(restype),
                field = ' ',
                resseq = resseq_res.item(),
                icode = icode_res,
            )

            for i, atom_name in enumerate(restype_to_heavyatom_names[restype]):
                if atom_name == '': continue    # No expected atom
                if (~mask_allatom_res[i]).any(): continue     # Atom is missing
                if len(atom_name) == 1: fullname = ' %s  ' % atom_name
                elif len(atom_name) == 2: fullname = ' %s ' % atom_name
                elif len(atom_name) == 3: fullname = ' %s' % atom_name
                else: fullname = atom_name # len == 4
                builder.init_atom(atom_name, pos_allatom_res[i].tolist(), 0.0, 1.0, ' ', fullname,)

    warnings.simplefilter('ignore', BiopythonWarning)
    builder = StructureBuilder()
    builder.init_structure(0)
    builder.init_model(0)

    unique_chain_nb = data['chain_nb'].unique().tolist()
    for ch_nb in unique_chain_nb:
        mask = (data['chain_nb'] == ch_nb)
        aa = _mask_select(data['aa'], mask)
        pos_heavyatom = _mask_select(data['pos_heavyatom'], mask)
        mask_heavyatom = _mask_select(data['mask_heavyatom'], mask)
        chain_id = _mask_select(data['chain_id'], mask)
        resseq = _mask_select(data['resseq'], mask)
        icode = _mask_select(data['icode'], mask)

        _build_chain(builder, aa, pos_heavyatom, mask_heavyatom, chain_id, resseq, icode)
    
    structure = builder.get_structure()
    if path is not None:
        io = PDBIO()
        io.set_structure(structure)
        io.save(path)
    return structure

AbDock/src/utils/protein/writers.py
pengzhangzhi/ab_opt
import numpy as np
import torch
from easydict import EasyDict

from .misc import BlackHole


def get_optimizer(cfg, model):
    if cfg.type == 'adam':
        return torch.optim.Adam(
            model.parameters(),
            lr=cfg.lr,
            weight_decay=cfg.weight_decay,
            betas=(cfg.beta1, cfg.beta2, )
        )
    else:
        raise NotImplementedError('Optimizer not supported: %s' % cfg.type)


def get_scheduler(cfg, optimizer):
    if cfg.type is None:
        return BlackHole()
    elif cfg.type == 'plateau':
        return torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            factor=cfg.factor,
            patience=cfg.patience,
            min_lr=cfg.min_lr,
        )
    elif cfg.type == 'multistep':
        return torch.optim.lr_scheduler.MultiStepLR(
            optimizer,
            milestones=cfg.milestones,
            gamma=cfg.gamma,
        )
    elif cfg.type == 'exp':
        return torch.optim.lr_scheduler.ExponentialLR(
            optimizer,
            gamma=cfg.gamma,
        )
    elif cfg.type is None:
        return BlackHole()
    else:
        raise NotImplementedError('Scheduler not supported: %s' % cfg.type)


def get_warmup_sched(cfg, optimizer):
    if cfg is None: return BlackHole()
    lambdas = [lambda it : (it / cfg.max_iters) if it <= cfg.max_iters else 1 for _ in optimizer.param_groups]
    warmup_sched = torch.optim.lr_scheduler.LambdaLR(optimizer, lambdas)
    return warmup_sched


def log_losses(out, it, tag, logger=BlackHole(), writer=BlackHole(), others={}):
    logstr = '[%s] Iter %05d' % (tag, it)
    logstr += ' | loss %.4f' % out['overall'].item()
    for k, v in out.items():
        if k == 'overall': continue
        logstr += ' | loss(%s) %.4f' % (k, v.item())
    for k, v in others.items():
       logstr += ' | %s %2.4f' % (k, v)
    logger.info(logstr)

    for k, v in out.items():
        if k == 'overall':
            writer.add_scalar('%s/loss' % tag, v, it)
        else:
            writer.add_scalar('%s/loss_%s' % (tag, k), v, it)
    for k, v in others.items():
        writer.add_scalar('%s/%s' % (tag, k), v, it)
    writer.flush()


class ValidationLossTape(object):

    def __init__(self):
        super().__init__()
        self.accumulate = {}
        self.others = {}
        self.total = 0

    def update(self, out, n, others={}):
        self.total += n
        for k, v in out.items():
            if k not in self.accumulate:
                self.accumulate[k] = v.clone().detach()
            else:
                self.accumulate[k] += v.clone().detach()

        for k, v in others.items():
            if k not in self.others:
                self.others[k] = v.clone().detach()
            else:
                self.others[k] += v.clone().detach()
        

    def log(self, it, logger=BlackHole(), writer=BlackHole(), tag='val'):
        avg = EasyDict({k:v / self.total for k, v in self.accumulate.items()})
        avg_others = EasyDict({k:v / self.total for k, v in self.others.items()})
        log_losses(avg, it, tag, logger, writer, others=avg_others)
        return avg['overall']


def recursive_to(obj, device):
    if isinstance(obj, torch.Tensor):
        if device == 'cpu':
            return obj.cpu()
        try:
            return obj.cuda(device=device, non_blocking=True)
        except RuntimeError:
            return obj.to(device)
    elif isinstance(obj, list):
        return [recursive_to(o, device=device) for o in obj]
    elif isinstance(obj, tuple):
        return tuple(recursive_to(o, device=device) for o in obj)
    elif isinstance(obj, dict):
        return {k: recursive_to(v, device=device) for k, v in obj.items()}

    else:
        return obj


def reweight_loss_by_sequence_length(length, max_length, mode='sqrt'):
    if mode == 'sqrt':
        w = np.sqrt(length / max_length)
    elif mode == 'linear':
        w = length / max_length
    elif mode is None:
        w = 1.0
    else:
        raise ValueError('Unknown reweighting mode: %s' % mode)
    return w


def sum_weighted_losses(losses, weights):
    """
    Args:
        losses:     Dict of scalar tensors.
        weights:    Dict of weights.
    """
    loss = 0
    for k in losses.keys():
        if weights is None:
            loss = loss + losses[k]
        else:
            loss = loss + weights[k] * losses[k]
    return loss


def count_parameters(model):
    return sum(p.numel() for p in model.parameters())

AbDock/src/utils/train.py
pengzhangzhi/ab_opt
# Transforms
from .mask import MaskSingleCDR, MaskMultipleCDRs, MaskAntibody, MaskFullAntibody
from .merge import MergeChains
from .patch import PatchAroundAnchor

# Factory
from ._base import get_transform, Compose

AbDock/src/utils/transforms/__init__.py
pengzhangzhi/ab_opt
import copy
import torch
# from torchvision.transforms import Compose


_TRANSFORM_DICT = {}

class Compose:
    """Composes several transforms together. This transform does not support torchscript.
    Please, see the note below.

    Args:
        transforms (list of ``Transform`` objects): list of transforms to compose.

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.PILToTensor(),
        >>>     transforms.ConvertImageDtype(torch.float),
        >>> ])

    .. note::
        In order to script the transformations, please use ``torch.nn.Sequential`` as below.

        >>> transforms = torch.nn.Sequential(
        >>>     transforms.CenterCrop(10),
        >>>     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        >>> )
        >>> scripted_transforms = torch.jit.script(transforms)

        Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require
        `lambda` functions or ``PIL.Image``.

    """

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img

    def __repr__(self) -> str:
        format_string = self.__class__.__name__ + "("
        for t in self.transforms:
            format_string += "\n"
            format_string += f"    {t}"
        format_string += "\n)"
        return format_string
    
def register_transform(name):
    def decorator(cls):
        _TRANSFORM_DICT[name] = cls
        return cls
    return decorator


def get_transform(cfg):
    if cfg is None or len(cfg) == 0:
        return None
    tfms = []
    for t_dict in cfg:
        t_dict = copy.deepcopy(t_dict)
        cls = _TRANSFORM_DICT[t_dict.pop('type')]
        tfms.append(cls(**t_dict))
    return Compose(tfms)


def _index_select(v, index, n):
    if isinstance(v, torch.Tensor) and v.size(0) == n:
        return v[index]
    elif isinstance(v, list) and len(v) == n:
        return [v[i] for i in index]
    else:
        return v


def _index_select_data(data, index):
    return {
        k: _index_select(v, index, data['aa'].size(0))
        for k, v in data.items()
    }


def _mask_select(v, mask):
    if isinstance(v, torch.Tensor) and v.size(0) == mask.size(0):
        return v[mask]
    elif isinstance(v, list) and len(v) == mask.size(0):
        return [v[i] for i, b in enumerate(mask) if b]
    else:
        return v


def _mask_select_data(data, mask):
    return {
        k: _mask_select(v, mask)
        for k, v in data.items()
    }

AbDock/src/utils/transforms/_base.py
pengzhangzhi/ab_opt
import torch
import random
from typing import List, Optional

from ..protein import constants
from ._base import register_transform


def random_shrink_extend(flag, min_length=5, shrink_limit=1, extend_limit=2):
    first, last = continuous_flag_to_range(flag)
    length = flag.sum().item()
    if (length - 2*shrink_limit) < min_length:
        shrink_limit = 0
    first_ext = max(0, first-random.randint(-shrink_limit, extend_limit))
    last_ext = min(last+random.randint(-shrink_limit, extend_limit), flag.size(0)-1)
    flag_ext = flag.clone()
    flag_ext[first_ext : last_ext+1] = True
    return flag_ext


def continuous_flag_to_range(flag):
    first = (torch.arange(0, flag.size(0))[flag]).min().item()
    last = (torch.arange(0, flag.size(0))[flag]).max().item()
    return first, last

@register_transform('mask_single_cdr')
class MaskSingleCDR(object):

    def __init__(self, selection=None, augmentation=True):
        super().__init__()
        cdr_str_to_enum = {
            'H1': constants.CDR.H1,
            'H2': constants.CDR.H2,
            'H3': constants.CDR.H3,
            'L1': constants.CDR.L1,
            'L2': constants.CDR.L2,
            'L3': constants.CDR.L3,
            'H_CDR1': constants.CDR.H1,
            'H_CDR2': constants.CDR.H2,
            'H_CDR3': constants.CDR.H3,
            'L_CDR1': constants.CDR.L1,
            'L_CDR2': constants.CDR.L2,
            'L_CDR3': constants.CDR.L3,
            'CDR3': 'CDR3',     # H3 first, then fallback to L3
        }
        assert selection is None or selection in cdr_str_to_enum
        self.selection = cdr_str_to_enum.get(selection, None)
        self.augmentation = augmentation

    def perform_masking_(self, data, selection=None):
        cdr_flag = data['cdr_flag']

        if selection is None:
            cdr_all = cdr_flag[cdr_flag > 0].unique().tolist()
            cdr_to_mask = random.choice(cdr_all)
        else:
            cdr_to_mask = selection

        cdr_to_mask_flag = (cdr_flag == cdr_to_mask)
        if self.augmentation:
            cdr_to_mask_flag = random_shrink_extend(cdr_to_mask_flag)

        cdr_first, cdr_last = continuous_flag_to_range(cdr_to_mask_flag)
        left_idx = max(0, cdr_first-1)
        right_idx = min(data['aa'].size(0)-1, cdr_last+1)
        anchor_flag = torch.zeros(data['aa'].shape, dtype=torch.bool)
        anchor_flag[left_idx] = True
        anchor_flag[right_idx] = True

        data['generate_flag'] = cdr_to_mask_flag
        data['anchor_flag'] = anchor_flag

    def __call__(self, structure):
        if self.selection is None:
            ab_data = []
            if structure['heavy'] is not None:
                ab_data.append(structure['heavy'])
            if structure['light'] is not None:
                ab_data.append(structure['light'])
            data_to_mask = random.choice(ab_data)
            sel = None
        elif self.selection in (constants.CDR.H1, constants.CDR.H2, constants.CDR.H3, ):
            data_to_mask = structure['heavy']
            sel = int(self.selection)
        elif self.selection in (constants.CDR.L1, constants.CDR.L2, constants.CDR.L3, ):
            data_to_mask = structure['light']
            sel = int(self.selection)
        elif self.selection == 'CDR3':
            if structure['heavy'] is not None:
                data_to_mask = structure['heavy']
                sel = constants.CDR.H3
            else:
                data_to_mask = structure['light']
                sel = constants.CDR.L3

        self.perform_masking_(data_to_mask, selection=sel)
        return structure


@register_transform('mask_multiple_cdrs')
class MaskMultipleCDRs(object):
    
    def __init__(self, selection: Optional[List[str]]=None, augmentation=True):
        super().__init__()
        cdr_str_to_enum = {
            'H1': constants.CDR.H1,
            'H2': constants.CDR.H2,
            'H3': constants.CDR.H3,
            'L1': constants.CDR.L1,
            'L2': constants.CDR.L2,
            'L3': constants.CDR.L3,
            'H_CDR1': constants.CDR.H1,
            'H_CDR2': constants.CDR.H2,
            'H_CDR3': constants.CDR.H3,
            'L_CDR1': constants.CDR.L1,
            'L_CDR2': constants.CDR.L2,
            'L_CDR3': constants.CDR.L3,
        }
        if selection is not None:
            self.selection = [cdr_str_to_enum[s] for s in selection]
        else:
            self.selection = None
        self.augmentation = augmentation

    def mask_one_cdr_(self, data, cdr_to_mask):
        cdr_flag = data['cdr_flag']

        cdr_to_mask_flag = (cdr_flag == cdr_to_mask)
        if self.augmentation:
            cdr_to_mask_flag = random_shrink_extend(cdr_to_mask_flag)

        cdr_first, cdr_last = continuous_flag_to_range(cdr_to_mask_flag)
        left_idx = max(0, cdr_first-1)
        right_idx = min(data['aa'].size(0)-1, cdr_last+1)
        anchor_flag = torch.zeros(data['aa'].shape, dtype=torch.bool)
        anchor_flag[left_idx] = True
        anchor_flag[right_idx] = True

        if 'generate_flag' not in data:
            data['generate_flag'] = cdr_to_mask_flag
            data['anchor_flag'] = anchor_flag
        else:
            data['generate_flag'] |= cdr_to_mask_flag
            data['anchor_flag'] |= anchor_flag

    def mask_for_one_chain_(self, data):
        cdr_flag = data['cdr_flag']
        cdr_all = cdr_flag[cdr_flag > 0].unique().tolist()
    
        num_cdrs_to_mask = random.randint(1, len(cdr_all))

        if self.selection is not None:
            cdrs_to_mask = list(set(cdr_all).intersection(self.selection))
        else:
            random.shuffle(cdr_all)
            cdrs_to_mask = cdr_all[:num_cdrs_to_mask]

        for cdr_to_mask in cdrs_to_mask:
            self.mask_one_cdr_(data, cdr_to_mask)

    def __call__(self, structure):
        if structure['heavy'] is not None:
            self.mask_for_one_chain_(structure['heavy'])
        if structure['light'] is not None:
            self.mask_for_one_chain_(structure['light'])
        return structure

@register_transform('mask_full_antibody')
class MaskFullAntibody(object):
    def __init__(self, antibody_chains: Optional[List[str]]=['heavy']) -> None:
        self.antibody_chain2enum = {
            'heavy': constants.Fragment.Heavy,
            'light': constants.Fragment.Light,
        }
        self.selections = [self.antibody_chain2enum[chain] for chain in antibody_chains]
    
    def maskone_chain_(self, data):
        data['generate_flag'] = torch.ones(data['aa'].shape, dtype=torch.bool)
        anchor_flag = torch.zeros(data['aa'].shape, dtype=torch.bool)
        anchor_flag[0] = True
        anchor_flag[data['aa'].size(0)-1] = True
        data['anchor_flag'] = anchor_flag
        
    def __call__(self, structure):
        """
        for selected antibody chains:
            - set the whole chain to be generated.
            - set the first and last residue to be anchors.
        """
        if (structure['heavy'] is not None) and \
            (constants.Fragment.Heavy in self.selections):
            self.maskone_chain_(structure['heavy'])
            
        if (structure['light'] is not None) and \
            (constants.Fragment.Light in self.selections):
            self.maskone_chain_(structure['light'])
        return structure

@register_transform('mask_antibody')
class MaskAntibody(object):

    def mask_ab_chain_(self, data):
        data['generate_flag'] = torch.ones(data['aa'].shape, dtype=torch.bool)

    def __call__(self, structure):
        pos_ab_alpha = []
        if structure['heavy'] is not None:
            self.mask_ab_chain_(structure['heavy'])
            pos_ab_alpha.append(
                structure['heavy']['pos_heavyatom'][:, constants.BBHeavyAtom.CA]
            )
        if structure['light'] is not None:
            self.mask_ab_chain_(structure['light'])
            pos_ab_alpha.append(
                structure['light']['pos_heavyatom'][:, constants.BBHeavyAtom.CA]
            )
        pos_ab_alpha = torch.cat(pos_ab_alpha, dim=0)   # (L_Ab, 3)

        if structure['antigen'] is not None:
            pos_ag_alpha = structure['antigen']['pos_heavyatom'][:, constants.BBHeavyAtom.CA]
            ag_ab_dist = torch.cdist(pos_ag_alpha, pos_ab_alpha)    # (L_Ag, L_Ab)
            nn_ab_dist = ag_ab_dist.min(dim=1)[0]   # (L_Ag)
            contact_flag = (nn_ab_dist <= 6.0)      # (L_Ag)
            if contact_flag.sum().item() == 0:
                contact_flag[nn_ab_dist.argmin()] = True

            anchor_idx = torch.multinomial(contact_flag.float(), num_samples=1).item()
            anchor_flag = torch.zeros(structure['antigen']['aa'].shape, dtype=torch.bool)
            anchor_flag[anchor_idx] = True
            structure['antigen']['anchor_flag'] = anchor_flag
            structure['antigen']['contact_flag'] = contact_flag
        
        return structure


@register_transform('remove_antigen')
class RemoveAntigen:

    def __call__(self, structure):
        structure['antigen'] = None
        structure['antigen_seqmap'] = None
        return structure

AbDock/src/utils/transforms/mask.py
pengzhangzhi/ab_opt
import torch

from ..protein import constants
from ._base import register_transform


@register_transform('merge_chains')
class MergeChains(object):

    def __init__(self):
        super().__init__()

    def assign_chain_number_(self, data_list):
        chains = set()
        for data in data_list:
            chains.update(data['chain_id'])
        chains = {c: i for i, c in enumerate(chains)}

        for data in data_list:
            data['chain_nb'] = torch.LongTensor([
                chains[c] for c in data['chain_id']
            ])

    def _data_attr(self, data, name):
        if name in ('generate_flag', 'anchor_flag') and name not in data:
            return torch.zeros(data['aa'].shape, dtype=torch.bool)
        else:
            return data[name]

    def __call__(self, structure):
        data_list = []
        if structure['heavy'] is not None:
            structure['heavy']['fragment_type'] = torch.full_like(
                structure['heavy']['aa'],
                fill_value = constants.Fragment.Heavy,
            )
            data_list.append(structure['heavy'])

        if structure['light'] is not None:
            structure['light']['fragment_type'] = torch.full_like(
                structure['light']['aa'],
                fill_value = constants.Fragment.Light,
            )
            data_list.append(structure['light'])

        if structure['antigen'] is not None:
            structure['antigen']['fragment_type'] = torch.full_like(
                structure['antigen']['aa'],
                fill_value = constants.Fragment.Antigen,
            )
            structure['antigen']['cdr_flag'] = torch.zeros_like(
                structure['antigen']['aa'],
            )
            data_list.append(structure['antigen'])

        self.assign_chain_number_(data_list)

        list_props = {
            'chain_id': [],
            'icode': [],
        }
        tensor_props = {
            'chain_nb': [],
            'resseq': [],
            'res_nb': [],
            'aa': [],
            'pos_heavyatom': [],
            'mask_heavyatom': [],
            'generate_flag': [],
            'cdr_flag': [],
            'anchor_flag': [],
            'fragment_type': [],
        }

        for data in data_list:
            for k in list_props.keys():
                list_props[k].append(self._data_attr(data, k))
            for k in tensor_props.keys():
                tensor_props[k].append(self._data_attr(data, k))

        list_props = {k: sum(v, start=[]) for k, v in list_props.items()}
        tensor_props = {k: torch.cat(v, dim=0) for k, v in tensor_props.items()}
        data_out = {
            **list_props,
            **tensor_props,
        }
        return data_out


AbDock/src/utils/transforms/merge.py
pengzhangzhi/ab_opt
import torch

from ._base import _mask_select_data, register_transform
from ..protein import constants

class DISTVIOLATION(Exception):
    pass
@register_transform('patch_around_anchor')
class PatchAroundAnchor(object):

    def __init__(
        self,
        initial_patch_size=128,
        antigen_size=128,
        remove_anchor=False,
        dist_cutoff=None,
        crop_contiguous_antigen=False,
        contiguous_threshold=1e6,
        contiguous_ratio=0.0,
        ):
        super().__init__()
        self.initial_patch_size = int(initial_patch_size)
        self.antigen_size = int(antigen_size)
        self.remove_anchor = remove_anchor if isinstance(remove_anchor,bool) else remove_anchor.lower() not in ['false', '0']
        self.dist_cutoff = dist_cutoff
        self.crop_contiguous_antigen = crop_contiguous_antigen if isinstance(crop_contiguous_antigen,bool) else crop_contiguous_antigen.lower() not in ['false', '0']
        self.contiguous_threshold = int(contiguous_threshold)
        self.contiguous_ratio = float(contiguous_ratio)
        
    def _center(self, data, origin):
        origin = origin.reshape(1, 1, 3)
        data['pos_heavyatom'] -= origin # (L, A, 3)
        data['pos_heavyatom'] = data['pos_heavyatom'] * data['mask_heavyatom'][:, :, None]
        data['origin'] = origin.reshape(3)
        return data

    def __call__(self, data):
        anchor_flag = data['anchor_flag']   # (L,)
        anchor_points = data['pos_heavyatom'][anchor_flag, constants.BBHeavyAtom.CA]    # (n_anchors, 3)
        antigen_mask = (data['fragment_type'] == constants.Fragment.Antigen)
        antibody_mask = torch.logical_not(antigen_mask)

        if anchor_flag.sum().item() == 0:
            # Generating full antibody-Fv, no antigen given
            data_patch = _mask_select_data(
                data = data,
                mask = antibody_mask,
            )
            data_patch = self._center(
                data_patch,
                origin = data_patch['pos_heavyatom'][:, constants.BBHeavyAtom.CA].mean(dim=0)
            )
            return data_patch

        pos_alpha = data['pos_heavyatom'][:, constants.BBHeavyAtom.CA]  # (L, 3)
        cdr_points = data['pos_heavyatom'][data['generate_flag']][:, constants.BBHeavyAtom.CA]  # (cdr_len, 3)
        dist_anchor = torch.cdist(pos_alpha, cdr_points).min(dim=1)[0]    # (L, )
        initial_patch_idx = torch.topk(
            dist_anchor,
            k = min(self.initial_patch_size, dist_anchor.size(0)),
            largest=False,
        )[1]   # (initial_patch_size, )

        dist_anchor_antigen = dist_anchor.masked_fill(
            mask = antibody_mask, # Fill antibody with +inf
            value = float('+inf')
        )   # (L, )
        min_dist = dist_anchor_antigen.min().item()
        if (self.dist_cutoff is not None) and \
            (min_dist > 0) and \
            (min_dist > self.dist_cutoff):
            raise DISTVIOLATION('Distance between anchor and antigen is too large')
        antigen_patch_idx = torch.topk(
            dist_anchor_antigen, 
            k = min(self.antigen_size, antigen_mask.sum().item()), 
            largest=False, sorted=True
        )[1]    # (ag_size, )
        if self.crop_contiguous_antigen and antigen_patch_idx.shape[0] > 0:
            contiguous_ratio, antigen_patch_idx = get_contiguous_idx(antigen_patch_idx, self.contiguous_threshold)
            if contiguous_ratio < self.contiguous_ratio:
                raise DISTVIOLATION(f'Antigen is not contiguous, keep ratio: {contiguous_ratio}')
        patch_mask = data['generate_flag'].clone()
        if not self.remove_anchor:
            patch_mask = torch.logical_or(
                patch_mask,
                anchor_flag,
            ) 
        
        patch_mask[initial_patch_idx] = True
        patch_mask[antigen_patch_idx] = True

        patch_idx = torch.arange(0, patch_mask.shape[0])[patch_mask]

        data_patch = _mask_select_data(data, patch_mask)
        data_patch = self._center(
            data_patch,
            origin = anchor_points.mean(dim=0)
        )
        data_patch['patch_idx'] = patch_idx
        return data_patch


def get_contiguous_idx(idx_tns, theshold=3):
    """
    Check if the indices are contiguous
    Args:
        idx_tns (torch.Tensor): indices tensor, shape (N, )
    Returns:
        contiguous_ratio (float): the ratio of the contiguous indices
        idx_tns (torch.Tensor): the indices tensor after removing the non-contiguous indices
    """
    idx_tns = idx_tns.sort()[0]
    diff = idx_tns[1:] - idx_tns[:-1]
    contiguous_tensor = torch.cat([torch.tensor([True]), diff <= theshold])
    # find the index of the first False value
    if contiguous_tensor.sum().item() == len(contiguous_tensor):
        return 1, idx_tns
    
    first_false_idx = torch.where(contiguous_tensor == False)[0][0]
    last_false_idx = torch.where(contiguous_tensor == False)[0][-1]
    if len(idx_tns) - 1 - last_false_idx > first_false_idx - 0:
        contiguous_tensor[:last_false_idx] = False
    else:
        contiguous_tensor[first_false_idx:] = False
    contiguous_ratio = round(contiguous_tensor.sum().item()/ len(contiguous_tensor),2)
    return contiguous_ratio, idx_tns[contiguous_tensor]
AbDock/src/utils/transforms/patch.py
pengzhangzhi/ab_opt

from ._base import register_transform


@register_transform('select_atom')
class SelectAtom(object):

    def __init__(self, resolution):
        super().__init__()
        assert resolution in ('full', 'backbone')
        self.resolution = resolution

    def __call__(self, data):
        if self.resolution == 'full':
            data['pos_atoms'] = data['pos_heavyatom'][:, :]
            data['mask_atoms'] = data['mask_heavyatom'][:, :]
        elif self.resolution == 'backbone':
            data['pos_atoms'] = data['pos_heavyatom'][:, :5]
            data['mask_atoms'] = data['mask_heavyatom'][:, :5]
        return data

AbDock/src/utils/transforms/select_atom.py
pengzhangzhi/ab_opt
import os
import shutil
import argparse
import sys
import torch
from torch.nn.utils import clip_grad_norm_
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

from src.datasets import get_dataset
from src.models import get_model
from src.utils.misc import *
from src.utils.data import *
from src.utils.train import *


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='configs/train/dock_antibody.yml')
    parser.add_argument('--logdir', type=str, default='./logs')
    parser.add_argument('--log', action='store_true', default=False)
    parser.add_argument('--device', type=str, default='cuda')
    parser.add_argument('--num_workers', type=int, default=8)
    parser.add_argument('--tag', type=str, default='')
    parser.add_argument('--resume', type=str, default=None)
    parser.add_argument('--finetune', type=str, default=None)
    args = parser.parse_args()

    # Load configs
    config, config_name = load_config(args.config)
    seed_all(config.train.seed)
    writer = BlackHole()
    if os.environ.get('wandb',False):
        import wandb
        wandb.init(project='diffab_dock', config=dict(config), name=f"{args.tag}")
    # Logging
    if args.log == False:
        logger = get_logger('train', None)
    else:
        if args.resume:
            log_dir = os.path.dirname(os.path.dirname(args.resume))
        else:
            log_dir = get_new_log_dir(args.logdir, prefix=config_name, tag=args.tag)
        ckpt_dir = os.path.join(log_dir, 'checkpoints')
        if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir)
        logger = get_logger('train', log_dir)
        if not os.path.exists(os.path.join(log_dir, os.path.basename(args.config))):
            shutil.copyfile(args.config, os.path.join(log_dir, os.path.basename(args.config)))
    logger.info(args)
    logger.info(config)

    # Data
    logger.info('Loading dataset...')
    train_dataset = get_dataset(config.dataset.train)
    val_dataset = get_dataset(config.dataset.val)
    train_iterator = inf_iterator(DataLoader(
        train_dataset, 
        batch_size=config.train.batch_size, 
        collate_fn=PaddingCollate(), 
        shuffle=True,
        num_workers=args.num_workers
    ))
    val_loader = DataLoader(val_dataset, batch_size=config.train.batch_size, collate_fn=PaddingCollate(), shuffle=False, num_workers=args.num_workers)
    logger.info('Train %d | Val %d' % (len(train_dataset), len(val_dataset)))

    # Model
    logger.info('Building model...')
    model = get_model(config.model).to(args.device)
    logger.info('Number of parameters: %d' % count_parameters(model))

    # Optimizer & scheduler
    optimizer = get_optimizer(config.train.optimizer, model)
    scheduler = get_scheduler(config.train.scheduler, optimizer)
    optimizer.zero_grad()
    it_first = 1

    # Resume
    is_resume = args.resume is not None
    is_finetune = args.finetune is not None
    if is_resume or is_finetune:
        ckpt_path = args.resume if is_resume else args.finetune
        logger.info('Resuming from checkpoint: %s' % ckpt_path)
        ckpt = torch.load(ckpt_path, map_location=args.device)
        it_first = ckpt['iteration']  if is_resume else 1
        msg = model.load_state_dict(ckpt['model'], strict=is_resume)
        logger.info(msg)
        if is_resume:
            logger.info('Resuming optimizer states...')
            optimizer.load_state_dict(ckpt['optimizer'])
            logger.info('Resuming scheduler states...')
            scheduler.load_state_dict(ckpt['scheduler'])

    # Train
    def train(it):
        time_start = current_milli_time()
        model.train()

        # Prepare data
        batch = recursive_to(next(train_iterator), args.device)

        # Forward
        # if args.debug: torch.set_anomaly_enabled(True)
        loss_dict = model(batch)
        loss = sum_weighted_losses(loss_dict, config.train.loss_weights)
        loss_dict['overall'] = loss
        time_forward_end = current_milli_time()

        # Backward
        loss.backward()
        orig_grad_norm = clip_grad_norm_(model.parameters(), config.train.max_grad_norm)
        optimizer.step()
        optimizer.zero_grad()
        time_backward_end = current_milli_time()
        if os.environ.get("wandb", False):
            wandb.log(
                {k:float(v) for k,v in loss_dict.items()}
            )
        # Logging
        log_losses(loss_dict, it, 'train', logger, writer, others={
            'grad': orig_grad_norm,
            'lr': optimizer.param_groups[0]['lr'],
            'time_forward': (time_forward_end - time_start) / 1000,
            'time_backward': (time_backward_end - time_forward_end) / 1000,
        })

        if not torch.isfinite(loss):
            logger.error('NaN or Inf detected.')
            torch.save({
                'config': config,
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'scheduler': scheduler.state_dict(),
                'iteration': it,
                'batch': recursive_to(batch, 'cpu'),
            }, os.path.join(log_dir, 'checkpoint_nan_%d.pt' % it))
            raise KeyboardInterrupt()

    # Validate
    def validate(it):
        loss_tape = ValidationLossTape()
        with torch.no_grad():
            model.eval()
            for i, batch in enumerate(tqdm(val_loader, desc='Validate', dynamic_ncols=True)):
                # Prepare data
                batch = recursive_to(batch, args.device)
                # Forward
                loss_dict = model(batch)
                loss = sum_weighted_losses(loss_dict, config.train.loss_weights)
                loss_dict['overall'] = loss

                loss_tape.update(loss_dict, 1)

        avg_loss = loss_tape.log(it, logger, writer, 'val')
        # Trigger scheduler
        if config.train.scheduler.type == 'plateau':
            scheduler.step(avg_loss)
        else:
            scheduler.step()
        return avg_loss

    try:
        for it in range(it_first, config.train.max_iters + 1):
            train(it)
            if it % config.train.val_freq == 0:
                avg_val_loss = validate(it)
                if args.log:
                    ckpt_path = os.path.join(ckpt_dir, '%d.pt' % it)
                    torch.save({
                        'config': config,
                        'model': model.state_dict(),
                        'optimizer': optimizer.state_dict(),
                        'scheduler': scheduler.state_dict(),
                        'iteration': it,
                        'avg_val_loss': avg_val_loss,
                    }, ckpt_path)
    except KeyboardInterrupt:
        logger.info('Terminating...')

AbDock/train.py
pengzhangzhi/ab_opt