peptidomimetics/objectives.py

import logging
import math
import os
import pickle
import sys
import warnings
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from lightning import pytorch as pl
from rdkit import Chem, RDLogger
from transformers import AutoModelWithLMHead, AutoTokenizer, EsmModel

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))
if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))

from BindEvaluator_models import *  # noqa: E402,F403
from DeepDTAGen_models import *  # noqa: E402,F403
from paths import peptiverse_root  # noqa: E402

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
torch.set_float32_matmul_precision("high")
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, message=r".*predict_dataloader.*many workers.*")
warnings.filterwarnings("ignore", message=r"Dropping last batch of size .*")
RDLogger.DisableLog("rdApp.*")
logging.basicConfig(level=logging.ERROR, force=True)
for name in ["lightning", "pytorch_lightning", "lightning.pytorch", "chemprop", "rdkit", "urllib3", "torch"]:
    logging.getLogger(name).setLevel(logging.ERROR)


_AMIDE_SMARTS = Chem.MolFromSmarts("[CX3](=[OX1])[NX3]")
_CARBONYL_C_SMARTS = Chem.MolFromSmarts("[CX3](=[OX1])")
_DIPEPTIDE_SMARTS = Chem.MolFromSmarts("[CX3](=[OX1])N[#6X4][CX3](=[OX1])N")

_PRED = None


def get_peptiverse_predictor():
    """Lazy PeptiVerse loader. Clone PeptiVerse and/or set PEPTIVERSE_ROOT first."""
    global _PRED
    if _PRED is not None:
        return _PRED

    root = peptiverse_root()
    if str(root) not in sys.path:
        sys.path.insert(0, str(root))
    from inference import PeptiVersePredictor

    device = "cuda" if torch.cuda.is_available() else "cpu"
    _PRED = PeptiVersePredictor(
        manifest_path=root / "best_models.txt",
        classifier_weight_root=root,
        device=device,
    )
    return _PRED


def _amide_bond_indices(mol, ignore_ring_amides=False):
    ids = set()
    for c_idx, _, n_idx in mol.GetSubstructMatches(_AMIDE_SMARTS):
        b = mol.GetBondBetweenAtoms(c_idx, n_idx)
        if b and b.GetBondType() == Chem.rdchem.BondType.SINGLE:
            if ignore_ring_amides and b.IsInRing():
                continue
            ids.add(b.GetIdx())
    return ids


def _carbonyl_c_indices(mol):
    return {m[0] for m in mol.GetSubstructMatches(_CARBONYL_C_SMARTS)}


def _carbonyl_neighbor_stats(mol, c_indices):
    stats = {"total": 0, "with_N": 0, "with_O": 0, "with_S": 0, "pure_amide": 0}
    for c_idx in c_indices:
        c = mol.GetAtomWithIdx(c_idx)
        stats["total"] += 1
        hasN = hasO = hasS = False
        for b in c.GetBonds():
            if b.GetBondType() != Chem.rdchem.BondType.SINGLE:
                continue
            z = b.GetOtherAtom(c).GetAtomicNum()
            if z == 7:
                hasN = True
            elif z == 8:
                hasO = True
            elif z == 16:
                hasS = True
        stats["with_N"] += int(hasN)
        stats["with_O"] += int(hasO)
        stats["with_S"] += int(hasS)
        if hasN and not (hasO or hasS):
            stats["pure_amide"] += 1
    return stats


def _adjacent_amide_pairs(mol):
    centers = set()
    for match in mol.GetSubstructMatches(_DIPEPTIDE_SMARTS):
        centers.add(match[3])
    return len(centers)


def analyze_peptide_likeness(smiles: str, ignore_ring_amides: bool = False, amide_density_target: float = 0.12):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Invalid SMILES: {smiles}")

    n_heavy_atoms = mol.GetNumAtoms()
    n_heavy_bonds = mol.GetNumBonds()
    amide_bonds = _amide_bond_indices(mol, ignore_ring_amides=ignore_ring_amides)
    n_amide_bonds = len(amide_bonds)
    carbonyl_cs = _carbonyl_c_indices(mol)
    cstats = _carbonyl_neighbor_stats(mol, carbonyl_cs)
    total_carb = max(1, cstats["total"])

    f1 = cstats["with_N"] / total_carb
    amide_per_atom = n_amide_bonds / max(1, n_heavy_atoms)
    f2 = min(1.0, amide_per_atom / max(1e-8, amide_density_target))
    f3 = 1.0 - math.exp(-_adjacent_amide_pairs(mol))
    penalty = 1.0 - (cstats["pure_amide"] / total_carb)
    score = max(0.0, min(1.0, 0.55 * f1 + 0.25 * f2 + 0.20 * f3 - 0.25 * penalty))

    return {
        "n_heavy_atoms": n_heavy_atoms,
        "n_heavy_bonds": n_heavy_bonds,
        "n_carbonyls": cstats["total"],
        "n_amide_bonds": n_amide_bonds,
        "amide_bond_ratio_all_bonds": n_amide_bonds / max(1, n_heavy_bonds),
        "acyl_N_fraction": f1,
        "pure_amide_fraction": cstats["pure_amide"] / total_carb,
        "amide_per_atom": amide_per_atom,
        "n_adjacent_amide_pairs": _adjacent_amide_pairs(mol),
        "peptide_likeness": score,
    }


class Toxicity:
    """PeptiVerse SMILES toxicity. Returns non-toxicity so mix_score stays unchanged."""

    def __call__(self, smiles_tokens, smiles_seqs):
        pred = get_peptiverse_predictor()
        scores = []
        for seq in smiles_seqs:
            score = pred.predict_property("toxicity", col="smiles", input_str=seq)["score"]
            scores.append(1.0 - score)
        return "non_toxicity", scores


class Solubility:
    def __call__(self, smiles_tokens, smiles_seqs):
        pred = get_peptiverse_predictor()
        scores = [pred.predict_property("solubility", col="smiles", input_str=seq)["score"] for seq in smiles_seqs]
        return "solubility", scores


class Permeability:
    def __call__(self, smiles_tokens, smiles_seqs):
        pred = get_peptiverse_predictor()
        scores = []
        for seq in smiles_seqs:
            score = pred.predict_property("permeability_pampa", col="smiles", input_str=seq)["score"]
            scores.append((score + 9.0) / (-4.0 + 9.0))
        return "permeability", scores


class Halflife:
    def __call__(self, smiles_tokens, smiles_seqs):
        pred = get_peptiverse_predictor()
        scores = [pred.predict_property("halflife", col="smiles", input_str=seq)["score"] for seq in smiles_seqs]
        return "halflife", scores


class BindingAffinity:
    def __init__(self, prot_seq, device=None):
        self.prot_seq = prot_seq
        self.device = device

    def __call__(self, smiles_tokens, smiles_seqs):
        pred = get_peptiverse_predictor()
        scores = []
        for seq in smiles_seqs:
            score = pred.predict_binding_affinity(col="smiles", target_seq=self.prot_seq, binder_str=seq)["affinity"]
            scores.append(min(10.0, score) / 10.0)
        return "affinity", scores


class SmallMolecule:
    def __init__(self, protein_sequence, device):
        import chemprop

        self.chemprop = chemprop
        self.trainer = pl.Trainer(logger=False, enable_progress_bar=False, accelerator="cuda" if torch.cuda.is_available() else "cpu", devices=1)
        ckpt_dir = REPO_ROOT / "peptidomimetics" / "ckpt" / "admetica"
        self.models = self.load_models(ckpt_dir)

        model_path = REPO_ROOT / "peptidomimetics" / "ckpt" / "deepdtagen" / "deepdtagen_model_bindingdb.pth"
        tokenizer_path = REPO_ROOT / "peptidomimetics" / "ckpt" / "deepdtagen" / "bindingdb_tokenizer.pkl"
        with open(tokenizer_path, "rb") as f:
            tokenizer = pickle.load(f)

        self.deep_dta_gen = DeepDTAGen(tokenizer)
        self.deep_dta_gen.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
        self.deep_dta_gen.to(device)
        self.deep_dta_gen.eval()
        self.protein_sequence = protein_sequence
        self.device = device

    def load_models(self, ckpt_dir):
        chemprop = self.chemprop
        toxicity_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "ld50.ckpt"))
        solubility_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "solubility.ckpt"))
        permeability_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "caco2.ckpt"))
        halflife_model = chemprop.models.MPNN.load_from_checkpoint(os.path.join(ckpt_dir, "half-life.ckpt"))
        return toxicity_model, solubility_model, permeability_model, halflife_model

    def is_valid_smiles(self, smiles):
        try:
            return Chem.MolFromSmiles(smiles) is not None
        except Exception:
            return False

    def prediction(self, smiles_list, trainer, model):
        valid_smiles = [smi for smi in smiles_list if self.is_valid_smiles(smi)]
        valid_indices = [i for i, smi in enumerate(smiles_list) if self.is_valid_smiles(smi)]
        invalid_indices = [i for i in range(len(smiles_list)) if i not in valid_indices]
        if not valid_smiles:
            return np.full(len(smiles_list), "", dtype=object)

        test_data = [self.chemprop.data.MoleculeDatapoint.from_smi(smi) for smi in valid_smiles]
        featurizer = self.chemprop.featurizers.SimpleMoleculeMolGraphFeaturizer()
        test_dataset = self.chemprop.data.MoleculeDataset(test_data, featurizer=featurizer)
        test_loader = self.chemprop.data.build_dataloader(test_dataset, shuffle=False)
        with torch.no_grad():
            predictions = trainer.predict(model, test_loader)
        predictions = [pred.item() for batch in predictions for pred in batch]
        for index in invalid_indices:
            predictions.insert(index, "")
        return predictions

    def non_toxicity_from_log10mgkg(self, scores, lo=1.0, hi=4.0):
        res = []
        for score in scores:
            score = max(lo, min(hi, score))
            res.append((score - lo) / (hi - lo))
        return res

    def __call__(self, smiles_tokens, smiles_seqs):
        scores = [self.prediction(smiles_seqs, self.trainer, model) for model in self.models]
        non_toxicity = self.non_toxicity_from_log10mgkg(scores[0])
        solubility = [(max(-12, min(2, score)) + 12) / 14 for score in scores[1]]
        permeability = [(max(-8, min(-3, score)) + 8) / 5 for score in scores[2]]
        halflife = scores[3]
        ratio = [analyze_peptide_likeness(smiles)["peptide_likeness"] for smiles in smiles_seqs]

        protein_sequences = [self.protein_sequence] * len(smiles_seqs)
        batch = process_latent_batch(smiles_seqs, protein_sequences).to(self.device)
        with torch.no_grad():
            affinity = self.deep_dta_gen(batch).squeeze(-1)
        affinity = [min(10, score.item()) / 10 for score in affinity]

        return "small_molecule", {
            "non_toxicity": non_toxicity,
            "solubility": solubility,
            "permeability": permeability,
            "halflife": halflife,
            "affinity": affinity,
            "ratio": ratio,
        }


def parse_motifs(motif: str) -> list:
    parts = motif.split(",")
    result = []
    for part in parts:
        part = part.strip()
        if "-" in part:
            start, end = map(int, part.split("-"))
            result.extend(range(start, end + 1))
        else:
            result.append(int(part))
    print(f"Target Motifs: {result}")
    return torch.tensor(result)


class BindEvaluator(pl.LightningModule):
    def __init__(self, cfg):
        super(BindEvaluator, self).__init__()
        self.esm_model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D").eval()
        for param in self.esm_model.parameters():
            param.requires_grad = False
        self.chemberta_model = AutoModelWithLMHead.from_pretrained("seyonec/ChemBERTa_zinc250k_v2_40k").roberta.eval()
        for param in self.chemberta_model.parameters():
            param.requires_grad = False
        self.repeated_module = RepeatedModule(
            cfg.model.n_layers, cfg.model.d_model, cfg.model.d_hidden,
            cfg.model.n_head, cfg.model.d_k, cfg.model.d_v, cfg.model.d_inner,
            dropout=cfg.model.dropout,
        )
        self.final_attention_layer = MultiHeadAttentionSequence(
            cfg.model.n_head, cfg.model.d_model, cfg.model.d_k, cfg.model.d_v, dropout=cfg.model.dropout
        )
        self.final_ffn = FFN(cfg.model.d_model, cfg.model.d_inner, dropout=cfg.model.dropout)
        self.output_projection_prot = nn.Linear(cfg.model.d_model, 1)

    def forward(self, binder_tokens, target_tokens):
        peptide_sequence = self.chemberta_model(**binder_tokens).last_hidden_state
        protein_sequence = self.esm_model(**target_tokens).last_hidden_state
        binder_mask = binder_tokens["attention_mask"]
        target_mask = target_tokens["attention_mask"]
        prot_enc, sequence_enc, sequence_attention_list, prot_attention_list, \
            prot_seq_attention_list, seq_prot_attention_list = self.repeated_module(
                peptide_sequence, protein_sequence, peptide_mask=binder_mask, protein_mask=target_mask
            )
        prot_enc, final_prot_seq_attention = self.final_attention_layer(
            prot_enc, sequence_enc, sequence_enc,
            key_padding_mask=binder_mask, query_padding_mask=target_mask,
        )
        prot_enc = self.final_ffn(prot_enc, padding_mask=target_mask)
        prot_enc = self.output_projection_prot(prot_enc)
        return prot_enc


class MotifModel:
    def __init__(self, cfg, target, motifs, device, specificity):
        self.cfg = cfg
        self.threshold = 0.918
        self.device = device
        self.specificity = specificity
        self.chemberta_tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa_zinc250k_v2_40k")
        self.esm_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
        self.target = self.esm_tokenizer(target, return_tensors="pt").to(device)
        self.motifs = parse_motifs(motifs).to(device)
        ckpt = cfg.inference.ckpt
        if not os.path.isabs(ckpt):
            ckpt = str(REPO_ROOT / ckpt)
        self.bindevaluator = BindEvaluator.load_from_checkpoint(ckpt, cfg=cfg, map_location=device)

    def __call__(self, smiles_tokens, smiles_seqs):
        binder = self.chemberta_tokenizer(
            smiles_seqs, return_tensors="pt", padding=True, truncation=True, max_length=512
        ).to(self.device)
        batch_size = binder["input_ids"].shape[0]
        L = self.target["input_ids"].shape[1]
        target = {
            "input_ids": self.target["input_ids"].expand(batch_size, L),
            "attention_mask": self.target["attention_mask"].expand(batch_size, L),
        }
        prediction = self.bindevaluator(binder, target).squeeze(-1)
        probs = torch.sigmoid(prediction)
        motif_scores = probs[:, self.motifs].mean(dim=1)
        if self.specificity:
            non_motif_probs = probs[:, [i for i in range(probs.shape[1]) if i not in self.motifs]]
            mask = non_motif_probs >= self.threshold
            count = mask.sum(dim=-1)
            specificity = 1 - count / (L - 2)
            return "motif", (motif_scores.tolist(), specificity.tolist())
        return "motif", motif_scores.tolist()