pCoMole/peptide/CNP_check.py

from typing import Tuple, List, Optional
from rdkit import Chem
from rdkit.Chem import Draw, rdDepictor

class DisulfideBuildError(Exception):
    pass


def _clean_peptide_sequence(seq: str) -> str:
    seq = seq.strip().upper().replace(" ", "")
    if not seq:
        raise DisulfideBuildError("Empty sequence.")
    return seq


def _validate_two_cys(seq: str, require_last_cys: bool = False) -> List[int]:
    cys_positions = [i for i, aa in enumerate(seq) if aa == "C"]

    if len(cys_positions) != 2:
        raise DisulfideBuildError(
            f"Sequence must contain exactly two cysteines, but found {len(cys_positions)}."
        )

    if require_last_cys and seq[-1] != "C":
        raise DisulfideBuildError("Last residue must be Cys.")

    return cys_positions


def _get_residue_number_to_sg_atom_idx(mol: Chem.Mol) -> dict:
    residue_to_sg = {}

    for atom in mol.GetAtoms():
        info = atom.GetPDBResidueInfo()
        if info is None:
            continue

        atom_name = info.GetName().strip()
        residue_name = info.GetResidueName().strip()
        residue_number = info.GetResidueNumber()

        if residue_name == "CYS" and atom_name == "SG":
            residue_to_sg[residue_number] = atom.GetIdx()

    return residue_to_sg


def _find_ss_bond_idx(mol: Chem.Mol) -> Optional[int]:
    for bond in mol.GetBonds():
        a1 = bond.GetBeginAtom()
        a2 = bond.GetEndAtom()
        if a1.GetSymbol() == "S" and a2.GetSymbol() == "S":
            return bond.GetIdx()
    return None


def _bond_is_in_ring(mol: Chem.Mol, bond_idx: int) -> bool:
    return any(bond_idx in ring for ring in mol.GetRingInfo().BondRings())


def build_peptide_with_disulfide(
    seq: str,
    require_last_cys: bool = False,
    require_ring_closure: bool = True,
) -> Chem.Mol:
    seq = _clean_peptide_sequence(seq)
    cys_positions = _validate_two_cys(seq, require_last_cys=require_last_cys)

    mol = Chem.MolFromFASTA(seq)
    if mol is None:
        raise DisulfideBuildError("RDKit failed to build peptide from sequence.")

    residue_to_sg = _get_residue_number_to_sg_atom_idx(mol)

    res_i = cys_positions[0] + 1
    res_j = cys_positions[1] + 1

    if res_i not in residue_to_sg or res_j not in residue_to_sg:
        raise DisulfideBuildError("Could not locate one or both cysteine SG atoms.")

    sg_i = residue_to_sg[res_i]
    sg_j = residue_to_sg[res_j]

    if mol.GetBondBetweenAtoms(sg_i, sg_j) is None:
        rw = Chem.RWMol(mol)
        rw.AddBond(sg_i, sg_j, Chem.BondType.SINGLE)
        mol = rw.GetMol()

        try:
            Chem.SanitizeMol(mol)
        except Exception as e:
            raise DisulfideBuildError(f"Sanitization failed after adding S-S bond: {e}")

    ss_bond_idx = _find_ss_bond_idx(mol)
    if ss_bond_idx is None:
        raise DisulfideBuildError("No S-S bond found after construction.")

    if require_ring_closure and not _bond_is_in_ring(mol, ss_bond_idx):
        raise DisulfideBuildError("S-S bond exists but does not close a ring.")

    return mol


def peptide_to_smiles_and_png(
    seq: str,
    png_path: str = "peptide_disulfide.png",
    require_last_cys: bool = False,
):
    mol = build_peptide_with_disulfide(
        seq,
        require_last_cys=require_last_cys,
        require_ring_closure=True,
    )

    smiles = Chem.MolToSmiles(mol)

    rdDepictor.Compute2DCoords(mol)
    img = Draw.MolToImage(mol, size=(1200, 800))
    img.save(png_path)

    ss_bond_idx = _find_ss_bond_idx(mol)
    has_ring = _bond_is_in_ring(mol, ss_bond_idx) if ss_bond_idx is not None else False

    return {
        "sequence": seq,
        "smiles": smiles,
        "png_path": png_path,
        "has_ss_bond": ss_bond_idx is not None,
        "ss_bond_in_ring": has_ring,
    }

seq = "KRSKPCFGDGKLDRQC"

result = peptide_to_smiles_and_png(
    seq,
    png_path="cnp_like_disulfide.png",
    require_last_cys=True,
)

print("SMILES:")
print(result["smiles"])
print("PNG saved to:", result["png_path"])
print("Has S-S bond:", result["has_ss_bond"])
print("S-S bond closes ring:", result["ss_bond_in_ring"])