SBVS/model/utils.py

import os
import h5py
import torch
import pickle
import numpy as np
from tqdm import tqdm
from rdkit import Chem
from rdkit.Chem import AllChem
from graphein.protein.graphs import construct_graph
from graphein.protein.config import ProteinGraphConfig
from graphein.ml import GraphFormatConvertor
from graphein.protein.edges.distance import (add_peptide_bonds,
                                             add_hydrogen_bond_interactions,
                                             add_disulfide_interactions,
                                             add_ionic_interactions,
                                             add_aromatic_interactions,
                                             add_aromatic_sulphur_interactions,
                                             add_cation_pi_interactions
                                            )
from chemprop.nn_utils import initialize_weights
from sklearn.metrics import roc_auc_score, average_precision_score

from utils import get_metric_func
from model.models import GGAP_CPI, GGAP_CPI_joint, KANO_ESM, GGAP_CPI_ablation
from model.loss import CompositeLoss
from KANO_model.utils import build_optimizer, build_lr_scheduler, build_loss_func, MolGraph

def set_up_model(args, logger):
    assert args.mode in ['train', 'retrain', 'finetune', 'inference', 'baseline_inference']
    if args.ablation == 'none':
        if args.train_model == 'GGAP_CPI' and args.dataset_type != 'joint':
            model = GGAP_CPI(args,
                            classification=True, multiclass=False,
                            multitask=False, prompt=True).to(args.device)
        elif args.dataset_type == 'joint':
            model = GGAP_CPI_joint(args,
                            classification=True, multiclass=False,
                            multitask=True, prompt=True).to(args.device)
        elif args.train_model == 'KANO_ESM':
            model = KANO_ESM(args, 
                            classification=True, multiclass=False, 
                            multitask=False, prompt=True).to(args.device)
        initialize_weights(model)
    else:
        model = GGAP_CPI_ablation(args, 
                                   classification=True, multiclass=False,
                                   multitask=False, prompt=True).to(args.device)

    if args.checkpoint_path is not None and args.ablation in ['none', 'GCN', 'Attn', 'ESM']:
        model.molecule_encoder.load_state_dict(torch.load(args.checkpoint_path, map_location='cpu'), strict=False)
        logger.info('load KANO pretrained model') if args.print else None
    logger.info(f'model: {model}') if args.print else None

    args.init_lr = 0.0001
    # Optimizers
    optimizer = build_optimizer(model, args)
    logger.info(f'optimizer: {optimizer}') if args.print else None

    # Learning rate schedulers
    scheduler = build_lr_scheduler(optimizer, args)
    logger.info(f'scheduler: {scheduler}') if args.print else None

    # Loss function
    loss_func = CompositeLoss(args, args.loss_func_wt).to(args.device)
    logger.info(f'loss function: {loss_func}, loss weights: {args.loss_func_wt}') if args.print else None

    args.metric_func = get_metric_func(args)
    logger.info(f'metric function: {args.metric_func}') if args.print else None
    args.previous_epoch = 0
    
    if args.mode == 'finetune':
        pretrained_dict = torch.load(os.path.join(args.model_path,
                                                f"{args.train_model}_best_model.pt"),
                                    map_location="cpu")["state_dict"]
        model_dict = model.state_dict()
        filtered_dict = {}
        for k, v in pretrained_dict.items():
            if k in model_dict and model_dict[k].shape == v.shape:
                filtered_dict[k] = v
        model_dict.update(filtered_dict)
        model.load_state_dict(model_dict)
        
        # model.load_state_dict(torch.load(os.path.join(args.model_path,
        #                                 f'{args.train_model}_best_model.pt'), map_location='cpu')['state_dict'])
        logger.info(f'load model from {args.model_path} for finetuning') if args.print else None
    elif args.mode == 'retrain':
        try:
            pre_file = torch.load(args.save_model_path, map_location='cpu')
        except:
            pre_file = torch.load(args.save_model_path.split('.')[0] + '_ft.pt', map_location='cpu')
        model.load_state_dict(pre_file['state_dict'])
        logger.info(f'load model from {args.save_model_path} for retraining') if args.print else None
        optimizer.load_state_dict(pre_file['optimizer'])
        logger.info(f'optimizer: {optimizer}') if args.print else None
        logger.info(f'load optimizer from {args.save_model_path} for retraining') if args.print else None
        args.previous_epoch = pre_file['epoch']
        logger.info(f'retrain from epoch {args.previous_epoch}, { args.epochs - args.previous_epoch} lasting') if args.print else None
        scheduler.load_state_dict(pre_file['scheduler'])
        logger.info(f'load scheduler from {args.save_model_path} for retraining') if args.print else None
    elif args.mode in ['inference', 'baseline_infernce']:
        model.cpu()
        try:
            pre_file = torch.load(args.save_best_model_path, map_location='cpu')
        except:
            pre_file = torch.load(args.save_best_model_path.split('.')[0] + '_ft.pt', map_location='cpu')

        model.load_state_dict(pre_file['state_dict'])
        model.to(args.device)
        logger.info(f'load model from {args.save_best_model_path} for inference') if args.print else None
    
    return args, model, optimizer, scheduler, loss_func
        
       
def generate_siamse_smi(data, query_prot_ids, 
                        support_dataset, support_prot,
                        strategy='random', num=1):
    query_smiles, query_labels = np.array(data.smiles()).flatten(), np.array(data.targets()).flatten()
    support_smiles, support_labels = np.array(support_dataset.smiles()).flatten(), \
                                     np.array(support_dataset.targets()).flatten()
    query_prot_ids, support_prot = np.array(query_prot_ids), np.array(support_prot)

    uni_prot = np.unique(np.array(query_prot_ids))
    smiles, label, siam_smiles, siam_label = [], [], [], []
    for prot in tqdm(uni_prot, desc='Generating siamese pairs'):
        q_smiles, q_label = query_smiles[np.where(query_prot_ids == prot)[0]],\
                            query_labels[np.where(query_prot_ids == prot)[0]]
        s_smiles, s_label = support_smiles[support_prot == prot], support_labels[support_prot == prot]
        if strategy == 'random':
            siamse_idx = np.random.choice(len(s_smiles), num*len(s_smiles))
            smiles.extend(np.repeat(q_smiles, num))
            label.extend(np.repeat(q_label, num))
            siam_smiles.extend(s_smiles[siamse_idx])
            siam_label.extend(s_label[siamse_idx])
        elif strategy == 'full':
            smiles.extend(np.repeat(q_smiles, len(s_smiles)))
            label.extend(np.repeat(q_label, len(s_smiles)))
            siam_smiles.extend(np.repeat(s_smiles, len(q_smiles)))
            siam_label.extend(np.repeat(s_label, len(q_smiles)))
        elif strategy == 'TopN_Sim':
            if len(query_prot_ids) == len(support_prot):
                if (query_prot_ids == support_prot).all():
                    siamse_idx, _ = calculate_topk_similarity(q_smiles, s_smiles, top_k=num+1)
                    siamse_idx = siamse_idx[:, 1:].flatten()
            else:
                siamse_idx, _ = calculate_topk_similarity(q_smiles, s_smiles, top_k=num)
                siamse_idx = siamse_idx.flatten()
            smiles.extend(np.repeat(q_smiles, num))
            label.extend(np.repeat(q_label, num))
            siam_smiles.extend(s_smiles[siamse_idx])
            siam_label.extend(s_label[siamse_idx])

    assert len(smiles) == len(siam_smiles)
    return [np.array(smiles), np.array(label)], [np.array(siam_smiles), np.array(siam_label)]


def generate_protein_graph(prot_dict):
    new_edge_funcs = {"edge_construction_functions": [add_peptide_bonds,
                                                        # add_aromatic_interactions,
                                                        add_hydrogen_bond_interactions,
                                                        add_disulfide_interactions,
                                                        add_ionic_interactions,
                                                        add_aromatic_sulphur_interactions,
                                                        add_cation_pi_interactions]
                        }
    config = ProteinGraphConfig(**new_edge_funcs)
    convertor = GraphFormatConvertor(src_format="nx", dst_format="pyg")
    
    for uni_id in (prot_dict.keys()):
        uni = uni_id.split('_')[0]
        try:
            g = construct_graph(config=config, uniprot_id=uni, verbose=False)
            prot_dict[uni_id] = prot_dict[uni_id] + [convertor(g)]
        except:
            logger.info('No PDB ID, try using AlphaFold2 predicted structure')
            try:
                fp = download_alphafold_structure(uni, aligned_score=False)
                g = construct_graph(config=config, path=fp, verbose=False)
                prot_dict[uni_id] = prot_dict[uni_id] + [convertor(g)]
            except:
                logger.info('No AlphaFold2 predicted structure found!!')
                prot_dict[uni_id] = prot_dict[uni_id] + [None]
            pass
    return prot_dict


def tanimoto_similarity_matrix(fps1, fps2):
    fp_matrix1 = np.array(fps1)
    fp_matrix2 = np.array(fps2)
    
    dot_product = np.dot(fp_matrix1, fp_matrix2.T)
    
    norm_sq1 = np.sum(fp_matrix1, axis=1)
    norm_sq2 = np.sum(fp_matrix2, axis=1)
    tanimoto_sim = dot_product / (norm_sq1[:, None] + norm_sq2[None, :] - dot_product)
    
    return tanimoto_sim


def calculate_topk_similarity(smiles_list1, smiles_list2, top_k=1):
    """
    Calculate the Tanimoto Similarity between SMILES strings based on ECFP4 fingerprints
    Then, return the indexs with topK similarity
    """
    mols1 = [Chem.MolFromSmiles(smile) for smile in smiles_list1]
    mols2 = [Chem.MolFromSmiles(smile) for smile in smiles_list2]
    # Calculate ECFP4 fingerprints
    fps1 = [AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048) for mol in mols1]
    fps2 = [AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048) for mol in mols2]
    
    fps1_np = np.array([np.frombuffer(fp.ToBitString().encode(), 'u1') - ord('0') for fp in fps1])
    fps2_np = np.array([np.frombuffer(fp.ToBitString().encode(), 'u1') - ord('0') for fp in fps2])

    # Calculate the Tanimoto similarity
    similarity_matrix = tanimoto_similarity_matrix(fps1, fps2)

    # return indexs with topN similarity
    topk_indices = np.argsort(-similarity_matrix, axis=1)[:, :top_k]
    return topk_indices, similarity_matrix


def generate_label(label, data_category, args=None):
    if isinstance(label, torch.Tensor):
        label = label.cpu().numpy().flatten()

    # for regression label: only consider pKi and pKd data
    if args.mode != 'inference':
        reg_label = [label[i] if data_category[i] <= args.type_thre else 999 for i in range(len(label))]
    else:
        reg_label = [label[i] for i in range(len(label))]
    # for classification label: 
        # strong binder/biner/weak binder: onely consider affinity data (pKi, pKd, pIC50, pEC50, pPotency)
        # non-binder: consider all data 
    if args.dataset_type == 'joint':
        # cls_label = [[0, 0, 0, 1] if (label[i] <= 4 or data_category[i] > 5) else \
        #             [0, 0, 1, 0] if (label[i] <= 6) and data_category[i] <= 5 else \
        #             [0, 1, 0, 0] if (label[i] <= 8) and data_category[i] <= 5 else \
        #             [1, 0, 0, 0] if (label[i] > 8) and data_category[i] <= 5 else \
        #             [999, 999, 999, 999] for i in range(len(label))]
        cls_label = [[0] if (label[i] <= 6) and data_category[i] <= 5 else \
                    [1] if (label[i] > 6) and data_category[i] <= 5 else \
                    [999] for i in range(len(label))]
        # cls_label = [[0] if (label[i] <= 6) and data_category[i] <= 5 else \
        #             [1] if (label[i] > 8) and data_category[i] <= 5 else \
        #             [999] for i in range(len(label))]
    else:
        cls_label = [[999, 999, 999, 999] for i in range(len(label))]
    return reg_label, cls_label


def calc_multiclass(pred, label, func):
    """
    Calculate AUC, AUPR, and Acc for a multiclassification task.
    """
    # Ensure inputs are numpy arrays
    if not isinstance(pred, np.ndarray):
        pred = np.array(pred)
    if not isinstance(label, np.ndarray):
        label = np.array(label)
    
    # Number of classes
    num_classes = pred.shape[1]
    
    # Store results
    auc_scores = {}
    aupr_scores = {}

    # Calculate per-class AUC and AUPR
    for i in range(num_classes):
        # True labels and predicted scores for class `i`
        true_labels = label[:, i]
        pred_scores = pred[:, i]

        # Ensure no NaNs in true_labels
        if np.isnan(true_labels).any():
            raise ValueError(f"NaN detected in true labels for class {i}")

        # AUC and AUPR for current class
        try:
            auc = roc_auc_score(true_labels, pred_scores)
            aupr = average_precision_score(true_labels, pred_scores)
        except ValueError as e:
            # Handle cases where AUC/AUPR cannot be calculated (e.g., all labels are 0 or 1)
            auc = np.nan
            aupr = np.nan

        auc_scores[f'class_{i}'] = auc
        aupr_scores[f'class_{i}'] = aupr

    # Calculate macro-average metrics
    macro_auc = np.nanmean(list(auc_scores.values()))
    macro_aupr = np.nanmean(list(aupr_scores.values()))

    # calculate accuracy
    pred_labels = np.argmax(pred, axis=1)
    true_labels = np.argmax(label, axis=1)
    acc = np.mean(pred_labels == true_labels)
    return dict(zip(func, [[macro_auc], [macro_aupr], [acc]]))


def calc_binaryclass(pred, label, func):
    """
    Calculate AUC, AUPR, and Acc for a binary classification task.
    pred: 1D array-like of predicted probabilities for the positive class.
    label: 1D array-like of ground truth labels (0 or 1).
    func: A list/tuple indicating the metric names for the output (e.g., ["AUC","AUPR","ACC"]).
    """
    if not isinstance(pred, np.ndarray):
        pred = np.array(pred)
    if not isinstance(label, np.ndarray):
        label = np.array(label)
    if np.isnan(label).any():
        raise ValueError("NaN detected in true labels.")
    try:
        auc = roc_auc_score(label, pred)
    except ValueError:
        auc = np.nan
    try:
        aupr = average_precision_score(label, pred)
    except ValueError:
        aupr = np.nan
    pred_binary = (pred >= 0.5).astype(int)
    acc = np.mean(pred_binary == label)
    return dict(zip(func, [[auc], [aupr], [acc]]))