model/boltzgen/task/analyze/analyze_utils.py

import io
import warnings
from pathlib import Path
import random
from typing import List
import subprocess
import re
import biotite
import hydride
from sklearn.cluster import DBSCAN
from Bio import PDB
from biotite import structure
from Bio.Seq import Seq


from matplotlib import pyplot as plt
import matplotlib as mpl

mpl.rcParams["figure.max_open_warning"] = 100

from Bio import Align

from boltzgen.data.rmsd_computation import get_true_coordinates
from boltzgen.model.loss.diffusion import weighted_rigid_align
from boltzgen.task.predict.data_from_generated import collate

import torch
import numpy as np
import pandas as pd
from tqdm import tqdm

from boltzgen.data import const
from boltzgen.data.data import Structure, biotite_array_from_feat
from boltzgen.data.write.mmcif import to_mmcif
from boltzgen.model.loss.validation import factored_lddt_loss, compute_subset_rmsd

from biotite.structure.sasa import sasa
from biotite.structure.info import vdw_radius_single, vdw_radius_protor
import biotite.structure.io.pdbx as pdbx
import biotite.structure.io.pdb as pdbio

TARGET_ID_RE = re.compile(
    r"^(?:(?:sample\d+_|batch\d+_|rank\d+_)+)?([^_]+)(?:_[^_]+)*?(?:_(?:gen))*$"
)


def _load_stack(path):
    path = Path(path)
    suffix = path.suffix.lower()
    if suffix in {".cif", ".mmcif"}:
        cif_file = pdbx.CIFFile.read(str(path))
        stack = pdbx.get_structure(cif_file, use_author_fields=False)
    elif suffix in {".pdb", ".ent"}:
        pdb_file = pdbio.PDBFile.read(str(path))
        stack = pdbio.get_structure(pdb_file, model=None)
    else:
        raise ValueError(f"Unsupported structure file extension: {suffix}")

    return stack


def compute_rmsd(atom_coords: torch.Tensor, pred_atom_coords: torch.Tensor):
    rmsd, _ = compute_subset_rmsd(
        atom_coords,
        pred_atom_coords,
        atom_mask=torch.ones_like(atom_coords[..., 0]),
        align_weights=torch.ones_like(atom_coords[..., 0]),
        subset_mask=torch.ones_like(atom_coords[..., 0]),
        multiplicity=1,
    )
    return rmsd


def make_histogram(
    df,
    column_name: str,
):
    data = df[column_name].dropna()
    fig, ax = plt.subplots(figsize=(6, 4))

    ax.hist(data, bins=50, alpha=0.7, edgecolor="black", linewidth=0.5)
    ax.axvline(data.mean(), color="red", linestyle="dashed", linewidth=1)

    ax.set_title(
        f"{column_name.replace('_', ' ').capitalize()} Distribution", fontsize=12
    )
    ax.set_xlabel(column_name.replace("_", " "), fontsize=10)
    ax.set_ylabel("Count", fontsize=10)

    ax.grid(True, axis="y", linestyle="--", linewidth=0.5, alpha=0.6)
    ax.tick_params(axis="both", which="major", labelsize=8)

    plt.tight_layout()
    return fig


def get_best_folding_sample(folded):
    confidence = 0.8 * folded["design_to_target_iptm"] + 0.2 * folded["design_ptm"]
    best_idx = np.argmax(confidence)

    # TODO: remove the "if k in folded"
    best_sample = {
        k: folded[k][best_idx] for k in const.eval_keys_confidence if k in folded
    }
    best_sample["coords"] = folded["coords"][best_idx]
    return best_sample


def get_fold_metrics(
    feat,
    folded,
    compute_lddts=True,
    prefix="",
):
    batch = collate([feat])
    diffusion_samples = batch["coords"].shape[0]
    best_sample = get_best_folding_sample(folded)

    # Compute RMSDs
    rmsd_out = get_true_coordinates(
        batch=batch,
        out={"sample_atom_coords": torch.from_numpy(best_sample["coords"])},
        diffusion_samples=1,
        symmetry_correction=False,
        protein_lig_rmsd=True,
    )
    true_coords_resolved_mask = rmsd_out["true_coords_resolved_mask"]

    # Add to metrics dictionary
    metrics = {}
    metrics["rmsd"] = rmsd_out.get("rmsd").item()
    metrics["rmsd_design"] = rmsd_out.get("rmsd_design").item()
    metrics["rmsd_target"] = rmsd_out.get("rmsd_target").item()
    metrics["rmsd_design_target"] = rmsd_out.get("rmsd_design_target").item()
    metrics["target_aligned_rmsd_design"] = rmsd_out.get(
        "target_aligned_rmsd_design"
    ).item()
    metrics["rmsd<2.5"] = bool(metrics["rmsd"] <= 2.5)
    metrics["target_aligned<2.5"] = bool(metrics["target_aligned_rmsd_design"] <= 2.5)
    metrics["designability_rmsd_2"] = bool(metrics["rmsd_design"] <= 2.0)
    metrics["designability_rmsd_4"] = bool(metrics["rmsd_design"] <= 4.0)

    # Comput LDDTs
    if compute_lddts:
        all_lddt_dict, _ = factored_lddt_loss(
            feats=batch,
            atom_mask=true_coords_resolved_mask,
            true_atom_coords=batch["coords"],
            pred_atom_coords=torch.from_numpy(best_sample["coords"]),
            multiplicity=diffusion_samples,
            exclude_ions=False,
        )
        metrics.update({f"lddt_{k}": v.max().item() for k, v in all_lddt_dict.items()})
        metrics["designability_lddt_60"] = bool(metrics["lddt_intra_design"] >= 0.6)
        metrics["designability_lddt_65"] = bool(metrics["lddt_intra_design"] >= 0.65)
        metrics["designability_lddt_70"] = bool(metrics["lddt_intra_design"] >= 0.7)
        metrics["designability_lddt_75"] = bool(metrics["lddt_intra_design"] >= 0.75)
        metrics["designability_lddt_80"] = bool(metrics["lddt_intra_design"] >= 0.8)
        metrics["designability_lddt_85"] = bool(metrics["lddt_intra_design"] >= 0.85)
        metrics["designability_lddt_90"] = bool(metrics["lddt_intra_design"] >= 0.9)

    # metrics without prefix (backbone only is the same as all atom)
    # TODO: remove the "if k in best_sample"
    confs = {k: best_sample[k] for k in const.eval_keys_confidence if k in best_sample}
    confs["min_interaction_pae<1.5"] = bool(confs["min_interaction_pae"] <= 1.5)
    confs["min_interaction_pae<2"] = bool(confs["min_interaction_pae"] <= 2.0)
    confs["min_interaction_pae<2.5"] = bool(confs["min_interaction_pae"] <= 2.5)
    confs["min_interaction_pae<3"] = bool(confs["min_interaction_pae"] <= 3)
    confs["min_interaction_pae<4"] = bool(confs["min_interaction_pae"] <= 4)
    confs["min_interaction_pae<5"] = bool(confs["min_interaction_pae"] <= 5)
    confs["design_ptm>80"] = bool(confs["design_ptm"] >= 0.8)
    confs["design_ptm>75"] = bool(confs["design_ptm"] >= 0.75)
    confs["design_iptm>80"] = bool(confs["design_iptm"] >= 0.8)
    confs["design_iptm>70"] = bool(confs["design_iptm"] >= 0.7)
    confs["design_iptm>60"] = bool(confs["design_iptm"] >= 0.6)
    confs["design_iptm>50"] = bool(confs["design_iptm"] >= 0.5)

    prefixed_metrics = {f"{prefix}{k}": v for k, v in metrics.items()}
    prefixed_metrics.update(confs)
    return prefixed_metrics


def count_noncovalents(feat):
    metrics = {}
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        biotite_array = biotite_array_from_feat(feat)
        biotite_array, _ = hydride.add_hydrogen(biotite_array)
        hbond = biotite.structure.hbond(biotite_array)
    donor_idxs, acceptor_idxs = hbond[:, 0], hbond[:, 2]
    donor_design_hbonds = int(
        (
            biotite_array.is_design[donor_idxs]
            & ~biotite_array.is_chain_design[acceptor_idxs]
        ).sum()
    )
    acceptor_design_hbonds = int(
        (
            ~biotite_array.is_chain_design[donor_idxs]
            & biotite_array.is_design[acceptor_idxs]
        ).sum()
    )
    metrics["plip_hbonds"] = donor_design_hbonds + acceptor_design_hbonds

    # saltbridges
    pos_atoms = biotite_array[biotite_array.charge > 0]
    neg_atoms = biotite_array[biotite_array.charge < 0]
    if len(neg_atoms) > 0 and len(pos_atoms) > 0:
        pos_neg_distances = torch.cdist(
            torch.as_tensor(pos_atoms.coord), torch.as_tensor(neg_atoms.coord)
        )
        pos_idxs, neg_idxs = torch.where(
            (pos_neg_distances > 0.5) & (pos_neg_distances < 5.5)
        )
        # only keep the ones between design and non design
        pos_design_sb = int(
            (pos_atoms.is_design[pos_idxs] & ~neg_atoms.is_chain_design[neg_idxs]).sum()
        )
        neg_design_sb = int(
            (~pos_atoms.is_chain_design[pos_idxs] & neg_atoms.is_design[neg_idxs]).sum()
        )
        metrics["plip_saltbridge"] = pos_design_sb + neg_design_sb
    else:
        metrics["plip_saltbridge"] = 0
    return metrics


def tm_score(coords1, coords2):
    num_atoms1 = coords1.shape[0]
    num_atoms2 = coords2.shape[0]

    atom_array1 = structure.AtomArray(num_atoms1)
    atom_array1.coord = coords1.numpy()
    atom_array1.element = np.array(["C"] * num_atoms1)
    atom_array1.atom_name = np.array(["CA"] * num_atoms1)
    atom_array1.res_name = np.array(["ALA"] * num_atoms1)
    atom_array1.chain_id = np.array(["A"] * num_atoms1)
    atom_array1.res_id = np.arange(1, num_atoms1 + 1)

    atom_array2 = structure.AtomArray(num_atoms2)
    atom_array2.coord = coords2.numpy()
    atom_array2.element = np.array(["C"] * num_atoms2)
    atom_array2.atom_name = np.array(["CA"] * num_atoms2)
    atom_array2.res_name = np.array(["ALA"] * num_atoms2)
    atom_array2.chain_id = np.array(["A"] * num_atoms2)
    atom_array2.res_id = np.arange(1, num_atoms2 + 1)

    try:
        # This fails with a value error if the structures are too dissimilar. In that event, we return 0 as the TM-Score
        aligned, transform, fixed_indices, mobile_indices = (
            structure.superimpose_structural_homologs(
                atom_array1, atom_array2, max_iterations=25
            )
        )
        tm_align_fixed = structure.tm_score(
            atom_array1,
            aligned,
            fixed_indices,
            mobile_indices,
        )
    except:
        tm_align_fixed = 0

    tm_score_rmsd_aligned = 0
    if num_atoms1 == num_atoms2:
        coords1 = weighted_rigid_align(
            coords1.float()[None],
            coords2.float()[None],
            weights=torch.ones(len(coords1)).float()[None],
            mask=torch.ones(len(coords2))[None],
        ).squeeze()

        atom_array1 = structure.AtomArray(num_atoms1)
        atom_array1.coord = coords1.numpy()
        atom_array1.element = np.array(["C"] * num_atoms1)
        atom_array1.atom_name = np.array(["CA"] * num_atoms1)
        atom_array1.res_name = np.array(["ALA"] * num_atoms1)
        atom_array1.chain_id = np.array(["A"] * num_atoms1)
        atom_array1.res_id = np.arange(1, num_atoms1 + 1)

        atom_array2 = structure.AtomArray(num_atoms2)
        atom_array2.coord = coords2.numpy()
        atom_array2.element = np.array(["C"] * num_atoms2)
        atom_array2.atom_name = np.array(["CA"] * num_atoms2)
        atom_array2.res_name = np.array(["ALA"] * num_atoms2)
        atom_array2.chain_id = np.array(["A"] * num_atoms2)
        atom_array2.res_id = np.arange(1, num_atoms2 + 1)
        try:
            _, _, fixed_indices, mobile_indices = (
                structure.superimpose_structural_homologs(
                    atom_array1, atom_array2, max_iterations=25
                )
            )
            tm_score_rmsd_aligned = structure.tm_score(
                atom_array1,
                atom_array2,
                fixed_indices,
                mobile_indices,
            )
        except:
            pass

    return tm_score_rmsd_aligned, tm_align_fixed


def vendi_from_sim(mat):
    mat = mat + mat.T
    np.fill_diagonal(mat, 1.0)
    eigvals, _ = np.linalg.eigh(mat / len(mat))
    eigvals = np.clip(eigvals, 0.0, None)
    return np.exp(np.nansum(-(eigvals * np.log(eigvals))))


def vendi_scores(
    all_ca_coords: List[np.ndarray],
    all_metrics: list = None,
    fold_metrics: bool = False,
    diversity_subset: int = None,
    compute_lddts: bool = True,
    compute_iptms: bool = True,
    compute_min_int_paes: bool = True,
    backbone_fold_metrics: bool = False,
    allatom_fold_metrics: bool = True,
) -> float:
    if fold_metrics or diversity_subset is not None:
        assert all_metrics is not None
    if all_metrics is not None:
        assert len(all_ca_coords) == len(all_metrics)
    if diversity_subset is not None and diversity_subset < len(all_ca_coords):
        indices = random.sample(range(len(all_ca_coords)), diversity_subset)
        all_metrics = [all_metrics[i] for i in indices]
        all_ca_coords = [all_ca_coords[i] for i in indices]
    N = len(all_ca_coords)
    tm = np.zeros((N, N), dtype=np.float32)
    tm_fixed = np.zeros((N, N), dtype=np.float32)

    for i in tqdm(range(N), desc="Computing structure diversity."):
        for j in range(i + 1, N):
            tm_score_rmsd_aligned, tm_fixeds = tm_score(
                all_ca_coords[i], all_ca_coords[j]
            )
            tm[i, j] = tm_score_rmsd_aligned
            tm_fixed[i, j] = tm_fixeds

    scores = {
        "vendi_tm_fixed": vendi_from_sim(tm_fixed),
        "vendi_tm_align": vendi_from_sim(tm),
    }
    prefixes = []
    if allatom_fold_metrics:
        prefixes.append("")
    if backbone_fold_metrics:
        prefixes.append("bb_")
    for prefix in prefixes:
        mask_2 = np.array([m[f"{prefix}designability_rmsd_2"] for m in all_metrics])
        mask_4 = np.array([m[f"{prefix}designability_rmsd_4"] for m in all_metrics])
        mask_25 = np.array([m[f"{prefix}rmsd<2.5"] for m in all_metrics])
        mask_target_25 = np.array(
            [m[f"{prefix}target_aligned<2.5"] for m in all_metrics]
        )
        scores.update(
            {
                f"vendi_tm_{prefix}rmsd<2.5": vendi_from_sim(
                    tm_fixed[mask_25][:, mask_25]
                )
                if np.sum(mask_25) > 0
                else 0.0,
                f"vendi_tm_{prefix}rmsd_2": vendi_from_sim(tm_fixed[mask_2][:, mask_2])
                if np.sum(mask_2) > 0
                else 0.0,
                f"vendi_tm_{prefix}rmsd_4": vendi_from_sim(tm_fixed[mask_4][:, mask_4])
                if np.sum(mask_4) > 0
                else 0.0,
                f"vendi_tm_{prefix}target_aligned_rmsd<2.5": vendi_from_sim(
                    tm_fixed[mask_target_25][:, mask_target_25]
                )
                if np.sum(mask_target_25) > 0
                else 0.0,
            }
        )

        if compute_lddts:
            mask_60 = np.array(
                [m[f"{prefix}designability_lddt_60"] for m in all_metrics]
            )
            mask_65 = np.array(
                [m[f"{prefix}designability_lddt_65"] for m in all_metrics]
            )
            mask_70 = np.array(
                [m[f"{prefix}designability_lddt_70"] for m in all_metrics]
            )
            mask_75 = np.array(
                [m[f"{prefix}designability_lddt_75"] for m in all_metrics]
            )
            mask_80 = np.array(
                [m[f"{prefix}designability_lddt_80"] for m in all_metrics]
            )
            mask_85 = np.array(
                [m[f"{prefix}designability_lddt_85"] for m in all_metrics]
            )
            mask_90 = np.array(
                [m[f"{prefix}designability_lddt_90"] for m in all_metrics]
            )
            scores.update(
                {
                    f"vendi_tm_{prefix}lddt_60": vendi_from_sim(
                        tm_fixed[mask_60][:, mask_60]
                    )
                    if np.sum(mask_60) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}lddt_65": vendi_from_sim(
                        tm_fixed[mask_65][:, mask_65]
                    )
                    if np.sum(mask_65) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}lddt_70": vendi_from_sim(
                        tm_fixed[mask_70][:, mask_70]
                    )
                    if np.sum(mask_70) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}lddt_75": vendi_from_sim(
                        tm_fixed[mask_75][:, mask_75]
                    )
                    if np.sum(mask_75) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}lddt_80": vendi_from_sim(
                        tm_fixed[mask_80][:, mask_80]
                    )
                    if np.sum(mask_80) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}lddt_85": vendi_from_sim(
                        tm_fixed[mask_85][:, mask_85]
                    )
                    if np.sum(mask_85) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}lddt_90": vendi_from_sim(
                        tm_fixed[mask_90][:, mask_90]
                    )
                    if np.sum(mask_90) > 0
                    else 0.0,
                }
            )
        if compute_iptms:
            mask_80 = np.array([m[f"design_iptm>80"] for m in all_metrics])
            mask_70 = np.array([m[f"design_iptm>70"] for m in all_metrics])
            mask_60 = np.array([m[f"design_iptm>60"] for m in all_metrics])
            mask_50 = np.array([m[f"design_iptm>50"] for m in all_metrics])
            scores.update(
                {
                    f"vendi_tm_{prefix}design_iptm_80": vendi_from_sim(
                        tm_fixed[mask_80][:, mask_80]
                    )
                    if np.sum(mask_80) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}design_iptm_70": vendi_from_sim(
                        tm_fixed[mask_70][:, mask_70]
                    )
                    if np.sum(mask_70) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}design_iptm_60": vendi_from_sim(
                        tm_fixed[mask_60][:, mask_60]
                    )
                    if np.sum(mask_60) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}design_iptm_50": vendi_from_sim(
                        tm_fixed[mask_50][:, mask_50]
                    )
                    if np.sum(mask_50) > 0
                    else 0.0,
                }
            )
        if compute_min_int_paes:
            mask_15 = np.array([m[f"min_interaction_pae<1.5"] for m in all_metrics])
            mask_2 = np.array([m[f"min_interaction_pae<2"] for m in all_metrics])
            mask_25 = np.array([m[f"min_interaction_pae<2.5"] for m in all_metrics])
            mask_3 = np.array([m[f"min_interaction_pae<3"] for m in all_metrics])
            mask_4 = np.array([m[f"min_interaction_pae<4"] for m in all_metrics])
            mask_5 = np.array([m[f"min_interaction_pae<5"] for m in all_metrics])
            scores.update(
                {
                    f"vendi_tm_{prefix}min_interaction_pae_1.5": vendi_from_sim(
                        tm_fixed[mask_15][:, mask_15]
                    )
                    if np.sum(mask_15) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}min_interaction_pae_2": vendi_from_sim(
                        tm_fixed[mask_2][:, mask_2]
                    )
                    if np.sum(mask_2) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}min_interaction_pae_2.5": vendi_from_sim(
                        tm_fixed[mask_25][:, mask_25]
                    )
                    if np.sum(mask_25) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}min_interaction_pae_3": vendi_from_sim(
                        tm_fixed[mask_3][:, mask_3]
                    )
                    if np.sum(mask_3) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}min_interaction_pae_4": vendi_from_sim(
                        tm_fixed[mask_4][:, mask_4]
                    )
                    if np.sum(mask_4) > 0
                    else 0.0,
                    f"vendi_tm_{prefix}min_interaction_pae_5": vendi_from_sim(
                        tm_fixed[mask_5][:, mask_5]
                    )
                    if np.sum(mask_5) > 0
                    else 0.0,
                }
            )

    return scores


def vendi_sequences(all_seqs: List[np.ndarray], diversity_subset: int = None) -> float:
    if diversity_subset is not None and diversity_subset < len(all_seqs):
        all_seqs = random.sample(all_seqs, diversity_subset)

    N = len(all_seqs)
    sims = np.zeros((N, N), dtype=np.float32)
    aligner = Align.PairwiseAligner()
    for i in tqdm(range(N), desc="Computing sequence diversity."):
        for j in range(i + 1, N):
            seq1 = Seq(all_seqs[i])
            seq2 = Seq(all_seqs[j])
            alignments = aligner.align(seq1, seq2)

            similarity = alignments[0].score / max(len(seq1), len(seq2))
            sims[i, j] = similarity

    return {
        "vendi_seq_sim": vendi_from_sim(sims),
    }


def compute_novelty_foldseek(
    indir: Path,
    outdir: Path,
    reference_db: Path,
    files: List[str],
    foldseek_binary: str = "/data/rbg/users/hstark/foldseek/bin/foldseek",
) -> pd.DataFrame:
    if len(files) == 0:
        return np.nan

    aln_tsv = outdir / "aln.tsv"
    tmp_dir = outdir / "tmp"

    cmd = [
        foldseek_binary,
        "easy-search",
        str(indir),
        str(reference_db),
        str(aln_tsv),
        str(tmp_dir),
        "--format-output",
        "query,target,alntmscore,qtmscore,ttmscore",
        "--alignment-type",
        "1",
        "--exhaustive-search",
        "1",
    ]

    subprocess.run(cmd, check=True)

    df = pd.read_csv(
        aln_tsv,
        sep="\t",
        names=["query", "target", "alntmscore", "qtmscore", "ttmscore"],
    )
    df["tmscore"] = (df["qtmscore"] + df["ttmscore"]) / 2
    df = df.groupby("query").max().reset_index()
    queries = [Path(f).stem for f in files]
    df = df.set_index("query").reindex(queries, fill_value=0.0).reset_index()
    df_novelty = df[["query", "tmscore"]].rename(columns={"tmscore": "novelty"})
    return df_novelty


def _radius(res_name: str, atom_name: str, element: str) -> float:
    """
    ProtOr radius with element fallback.
    """
    try:
        r = vdw_radius_protor(res_name, atom_name)
        if r is not None:
            return r
    except KeyError:
        pass
    r = vdw_radius_single(element)
    return r if r is not None else 1.8


def compute_sasa(structure_path):
    HYDROPHOBIC_RESIDUES = {"ALA", "VAL", "LEU", "ILE", "MET", "PHE", "PRO", "TRP"}
    atoms = _load_stack(structure_path)[0]
    res_names = [
        bytes(r).decode() if isinstance(r, bytes) else str(r) for r in atoms.res_name
    ]
    atom_names = [
        bytes(a).decode() if isinstance(a, bytes) else str(a) for a in atoms.atom_name
    ]
    elements = [
        bytes(e).decode() if isinstance(e, bytes) else str(e) for e in atoms.element
    ]

    radii = np.array(
        [
            _radius(rn.strip(), an.strip(), el)
            for rn, an, el in zip(res_names, atom_names, elements)
        ],
        dtype=float,
    )
    atom_sasa = sasa(atoms, probe_radius=1.4, point_number=960, vdw_radii=radii)
    mask = (
        np.array([rn.strip() in HYDROPHOBIC_RESIDUES for rn in res_names])
        & np.char.startswith(atom_names, "C")
        & (atom_sasa > 0)
    )

    return atoms.coord[mask], atom_sasa[mask]


def largest_hydrophobic_patch_area(cif_path, distance_cutoff=6.0):
    result = compute_sasa(cif_path)
    if result is None or result[0].size == 0:
        return np.nan

    coords, sasa_vals = result
    if len(coords) == 0:
        return 0.0
    clustering = DBSCAN(eps=distance_cutoff, min_samples=1).fit(coords)
    labels = clustering.labels_

    max_patch_area = 0.0
    for label in np.unique(labels):
        area = sasa_vals[labels == label].sum()
        max_patch_area = max(max_patch_area, area)

    return max_patch_area


def get_delta_sasa(
    path,
    atom_target_mask,                
    atom_design_mask,           
):
    stack = _load_stack(path)
    atoms = stack[0]

    res = [
        r.decode().strip() if isinstance(r, bytes) else str(r).strip()
        for r in atoms.res_name
    ]
    atm = [
        a.decode().strip() if isinstance(a, bytes) else str(a).strip()
        for a in atoms.atom_name
    ]
    elem = [e.decode() if isinstance(e, bytes) else str(e) for e in atoms.element]

    radii = np.array(
        [_radius(rn, an, el) for rn, an, el in zip(res, atm, elem)], dtype=float
    )

    
    bound_mask = atom_design_mask | atom_target_mask
    atoms_bound = atoms[bound_mask]
    radii_bound = radii[bound_mask]

    area_bound = sasa(
        atoms_bound,
        probe_radius=1.4,
        point_number=960,
        vdw_radii=radii_bound,
    )
    
    target_in_bound = atom_target_mask[bound_mask]
    target_bound    = area_bound[target_in_bound].sum()
    
    

    target_atoms = atoms[atom_target_mask]
    target_res = [r for r, m in zip(res, atom_target_mask) if m]
    target_atm = [a for a, m in zip(atm, atom_target_mask) if m]
    target_elem = [e for e, m in zip(elem, atom_target_mask) if m]

    radii_lig = np.array(
        [_radius(rn, an, el) for rn, an, el in zip(target_res, target_atm, target_elem)],
        dtype=float,
    )
    target_area = sasa(
        target_atoms,
        probe_radius=1.4,
        point_number=960,
        vdw_radii=radii_lig,
    )
    delta = target_area.sum() - target_bound
    return delta, target_area.sum(), target_bound


def compute_ss_metrics(dssp_pred, ss_conditioning_metricsed):
    ss_metrics = {}
    conditioned_mask = ss_conditioning_metricsed != 0
    if conditioned_mask.sum() == 0:
        return {
            "precision_loop": float("nan"),
            "recall_loop": float("nan"),
            "accuracy_loop": float("nan"),
            "precision_helix": float("nan"),
            "recall_helix": float("nan"),
            "accuracy_helix": float("nan"),
            "precision_sheet": float("nan"),
            "recall_sheet": float("nan"),
            "accuracy_sheet": float("nan"),
            "accuracy_overall": float("nan"),
        }
    types = {1: "loop", 2: "helix", 3: "sheet"}
    TP_total, total_conditioned = 0, conditioned_mask.sum().item()
    for i, name in types.items():
        TP = ((dssp_pred == i) & (ss_conditioning_metricsed == i)).sum().item()
        FP = (
            ((dssp_pred == i) & (ss_conditioning_metricsed != i) & conditioned_mask)
            .sum()
            .item()
        )
        FN = ((dssp_pred != i) & (ss_conditioning_metricsed == i)).sum().item()
        precision = TP / (TP + FP) if (TP + FP) > 0 else float("nan")
        recall = TP / (TP + FN) if (TP + FN) > 0 else float("nan")
        accuracy = TP / (TP + FP + FN) if (TP + FP + FN) > 0 else float("nan")
        ss_metrics[f"precision_{name}"] = precision
        ss_metrics[f"recall_{name}"] = recall
        ss_metrics[f"accuracy_{name}"] = accuracy
        TP_total += TP

    accuracy_overall = (
        TP_total / total_conditioned if total_conditioned > 0 else float("nan")
    )
    ss_metrics["accuracy_overall"] = accuracy_overall
    return ss_metrics


def filter_resolved_atoms(structure: Structure) -> Structure:
    resolved_atom_indices = np.where(structure.atoms["is_present"])[0]
    return Structure.extract_atoms(structure, resolved_atom_indices)


def save_design_only_structure_to_cif(atom_design_mask, structure, output_path: Path):
    design_atom_indices = torch.where(atom_design_mask)[0].cpu().numpy()
    design_only_str = Structure.extract_atoms(
        structure, design_atom_indices, res_reindex=True
    )
    cif_text = to_mmcif(design_only_str)
    output_path.write_text(cif_text)
    return cif_text


def save_design_only_structure_to_pdb(atom_design_mask, structure, output_path: Path):
    cif_path = output_path.with_suffix(".cif")
    cif_text = save_design_only_structure_to_cif(atom_design_mask, structure, cif_path)

    cif_io = io.StringIO(cif_text)
    mmcif_parser = PDB.MMCIFParser()
    pdb_writer = PDB.PDBIO()
    parsed_structure = mmcif_parser.get_structure("des_only", cif_io)
    pdb_writer.set_structure(parsed_structure)
    pdb_writer.save(str(output_path))


########################################################################################################
# Hydrophobycity computation functions. From Jeremie Alexander.
########################################################################################################


def calc_base_h(seq: str) -> float:
    s = seq.upper()
    n = len(s)
    if n == 0:
        return 0.0
    H = 0.0

    # position-specific coefficients
    for i, aa in enumerate(s):
        if aa not in const.hydrophobicity_info:
            raise ValueError(f"Unknown residue '{aa}' in '{seq}'")
        if i == 0:
            key = "Rc1"
        elif i == 1:
            key = "Rc2"
        elif i == n - 1:
            key = "Rn"
        elif i == n - 2:
            key = "Rn1"
        else:
            key = "Rc"
        H += const.hydrophobicity_info[aa][key]

    # nearest-neighbor penalties around H/R/K
    for i, aa in enumerate(s):
        if aa in ("H", "R", "K"):
            for j in (i - 1, i + 1):
                if 0 <= j < n and s[j] in const.nn_penalty:
                    H -= const.nn_penalty[s[j]]

    # proline run penalties
    i = 0
    while i < n:
        if s[i] == "P":
            j = i
            while j < n and s[j] == "P":
                j += 1
            run = j - i
            if run >= 4:
                H -= 5.0
            elif run == 3:
                H -= 3.5
            elif run == 2:
                H -= 1.2
            i = j
        else:
            i += 1
    return H


def apply_length_weight(H: float, n: int) -> float:
    if n < 8:
        KL = 1.0 - 0.055 * (8 - n)
    elif n > 20:
        KL = 1.0 / (1.0 + 0.027 * (n - 20))
    else:
        KL = 1.0
    return H * KL


def overall_penalty(H: float) -> float:
    if H <= 20:
        return H
    if H <= 30:
        return H - 0.27 * (H - 18.0)
    if H <= 40:
        return H - 0.33 * (H - 18.0)
    if H <= 50:
        return H - 0.38 * (H - 18.0)
    return H - 0.447 * (H - 18.0)


def calc_hydrophobicity(seq: str) -> float:
    s = (seq or "").strip().upper()
    if not s or "X" in s:
        return float("nan")
    base = calc_base_h(s)
    base = apply_length_weight(base, len(s))
    return round(overall_penalty(base), 4)


########################################################################################################
# Liability computation functions
########################################################################################################


def get_motif_set(modality="antibody", peptide_type="linear"):
    """
    Returns a dict of motif_name -> compiled regex for the given modality.
    modality: 'antibody' or 'peptide'
    peptide_type: 'linear' or 'cyclic' (only for peptide modality)


    """
    if modality == "antibody":
        return {
            "DeAmdH": re.compile(r"N[GS]"),  # High-risk deamidation
            "DeAmdM": re.compile(r"N[AHNT]"),  # Medium-risk deamidation
            "DeAmdL": re.compile(r"[STK]N"),  # Low-risk deamidation
            "Ngly": re.compile(r"N[^P][ST]"),  # N-glycosylation sites
            "Isom": re.compile(r"D[DG HST]".replace(" ", "")),  # Isomerization
            "Isomer": re.compile(r"DG|DS|DD"),  # Isomerization variants
            "FragH": re.compile(r"DP"),  # High fragmentation risk
            "FragM": re.compile(r"TS"),  # Medium fragmentation risk
            "TrpOx": re.compile(r"W"),  # Tryptophan oxidation
            "MetOx": re.compile(r"M"),  # Methionine oxidation
            "Hydro": re.compile(r"NP"),  # Hydrolysis prone
            "IntBind": re.compile(r"GPR|RGD|RYD|LDV|DGE|KGD|NGR"),  # Integrin binding
            "Polyreactive": re.compile(
                r"GGG|GG|RR|VG|VVV|WWW|YY|WxW"
            ),  # Polyreactivity
            "AggPatch": re.compile(r"FHW"),  # Aggregation patches
            "ViscPatch": re.compile(r"HYF|HWH"),  # Viscosity patches
            "DeAmdH": re.compile(r"N[GS]"),
            "DeAmdM": re.compile(r"N[AHNT]"),
            "DeAmdL": re.compile(r"[STK]N"),
            "Ngly": re.compile(r"N[^P][ST]"),
            "Isom": re.compile(r"D[DG HST]".replace(" ", "")),
            "Isomer": re.compile(r"DG|DS|DD"),
            "FragH": re.compile(r"DP"),
            "FragM": re.compile(r"TS"),
            "TrpOx": re.compile(r"W"),
            "MetOx": re.compile(r"M"),
            "Hydro": re.compile(r"NP"),
            "IntBind": re.compile(r"GPR|RGD|RYD|LDV|DGE|KGD|NGR"),
            "Polyreactive": re.compile(r"GGG|GG|RR|VG|VVV|WWW|YY|WxW"),
            "AggPatch": re.compile(r"FHW"),
            "ViscPatch": re.compile(r"HYF|HWH"),
        }
    elif modality == "peptide":
        motifs = {
            "AspBridge": re.compile(r"N[GSQA]"),  # Deamidation hotspots
            "AspCleave": re.compile(r"D[PGS]"),  # Acidic cleavage sites
            "NTCycl": re.compile(r"^[QN]"),  # N-terminal cyclization
            "ProtTryp": re.compile(r"[KR](?=.)"),  # Trypsin cleavage sites
            "DPP4": re.compile(r"^[PX]?[AP]"),  # DPP4 cleavage sites
            "MetOx": re.compile(r"M"),  # Methionine oxidation
            "TrpOx": re.compile(r"W"),  # Tryptophan oxidation
            "HydroPatch": re.compile(r"[FILVWY]{3,}"),  # Hydrophobic patches
        }
        if peptide_type == "cyclic":
            # remove N-term liabilities for cyclic peptides
            motifs.pop("NTCycl", None)
            motifs.pop("DPP4", None)
        elif peptide_type == "linear":
            # For linear peptides, we handle cysteine pairing separately
            # so we don't include CysOx in the motif set
            pass
        return motifs
    else:
        raise ValueError(f"Unknown modality: {modality}")


def severity_score(name):
    return const.liability_severity.get(name, const.default_severity)


def compute_liability_scores(sequences, modality="antibody", peptide_type="linear"):
    """
    Compute liability scores for given sequences.
    modality: 'antibody' or 'peptide'; peptide_type: 'linear' or 'cyclic'.
    For cyclic peptides, terminal CysOx flags are skipped.

    Returns:
        dict: sequence -> {'score': int, 'violations': list of dicts}
    """
    motifs = get_motif_set(modality, peptide_type)
    results = {}
    for seq in sequences:
        violations = []
        total_score = 0
        length = len(seq)
        # motif scanning
        for name, pat in motifs.items():
            for m in pat.finditer(seq):
                pos = m.start() + 1
                # skip terminal cysteines for cyclic peptides
                if (
                    modality == "peptide"
                    and peptide_type == "cyclic"
                    and name == "CysOx"
                    and pos in (1, length)
                ):
                    continue
                sev = severity_score(name)
                violations.append(
                    {"motif": name, "pos": pos, "len": len(m.group()), "severity": sev}
                )
                total_score += sev
        # antibody-specific extras
        if modality == "antibody":
            # unpaired cysteines
            cpos = [i for i, aa in enumerate(seq) if aa == "C"]
            paired = set()
            for i in range(len(cpos) - 1):
                if abs(cpos[i + 1] - cpos[i]) in (1, 2):
                    paired.update({cpos[i], cpos[i + 1]})
            for i in cpos:
                if i not in paired:
                    sev = severity_score("UnpairedCys")
                    violations.append(
                        {
                            "motif": "UnpairedCys",
                            "pos": i + 1,
                            "len": 1,
                            "severity": sev,
                        }
                    )
                    total_score += sev
            # net charge
            charge = seq.count("K") + seq.count("R") - seq.count("D") - seq.count("E")
            if charge > 1:
                sev = const.default_severity
                violations.append(
                    {"motif": "HighNetCharge", "pos": None, "len": 0, "severity": sev}
                )
                total_score += sev
        # peptide-specific extras
        elif modality == "peptide":
            # For linear peptides, only flag unpaired cysteines (odd number of cysteines)
            # For cyclic peptides, terminal cysteines are expected, so only flag internal unpaired cysteines
            cpos = [i for i, aa in enumerate(seq) if aa == "C"]
            if peptide_type == "linear":
                # For linear peptides, if there's an odd number of cysteines, flag all cysteines as potential liabilities
                # since we don't know which one is unpaired
                if len(cpos) % 2 == 1:
                    sev = severity_score("UnpairedCys")
                    for cys_pos in cpos:
                        violations.append(
                            {
                                "motif": "UnpairedCys",
                                "pos": cys_pos + 1,
                                "len": 1,
                                "severity": sev,
                            }
                        )
                        total_score += sev
            elif peptide_type == "cyclic":
                # For cyclic peptides, terminal cysteines are expected for cyclization
                # Only flag internal unpaired cysteines
                internal_cpos = [
                    i for i in cpos if i != 0 and i != len(seq) - 1
                ]  # exclude terminal positions
                if len(internal_cpos) % 2 == 1:
                    # Flag all internal cysteines as potential liabilities
                    sev = severity_score("UnpairedCys")
                    for cys_pos in internal_cpos:
                        violations.append(
                            {
                                "motif": "UnpairedCys",
                                "pos": cys_pos + 1,
                                "len": 1,
                                "severity": sev,
                            }
                        )
                        total_score += sev

                # Additional liability checks for cyclic peptides
                # 1. Check for low hydrophilic content (< 40%)
                hydrophilic_residues = (
                    seq.count("D")
                    + seq.count("E")
                    + seq.count("K")
                    + seq.count("R")
                    + seq.count("H")
                    + seq.count("N")
                    + seq.count("Q")
                    + seq.count("S")
                    + seq.count("T")
                )
                hydrophilic_percentage = (hydrophilic_residues / len(seq)) * 100
                if hydrophilic_percentage < 40:
                    sev = severity_score("LowHydrophilic")
                    violations.append(
                        {
                            "motif": "LowHydrophilic",
                            "pos": None,
                            "len": 0,
                            "severity": sev,
                            "details": f"{hydrophilic_percentage:.1f}% hydrophilic",
                        }
                    )
                    total_score += sev

                # 2. Check for consecutive identical residues
                max_consec_identical = 1
                current_consec = 1
                for i in range(1, len(seq)):
                    if seq[i] == seq[i - 1]:
                        current_consec += 1
                        max_consec_identical = max(max_consec_identical, current_consec)
                    else:
                        current_consec = 1

                if max_consec_identical > 1:
                    sev = severity_score("ConsecIdentical")
                    violations.append(
                        {
                            "motif": "ConsecIdentical",
                            "pos": None,
                            "len": 0,
                            "severity": sev,
                            "details": f"{max_consec_identical} consecutive identical",
                        }
                    )
                    total_score += sev

                # 3. Check for more than 4 consecutive hydrophobic residues
                max_consec_hydrophobic = 0
                current_consec = 0
                for aa in seq:
                    if aa in "FILVWY":
                        current_consec += 1
                        max_consec_hydrophobic = max(
                            max_consec_hydrophobic, current_consec
                        )
                    else:
                        current_consec = 0

                if max_consec_hydrophobic > 4:
                    sev = severity_score("LongHydrophobic")
                    violations.append(
                        {
                            "motif": "LongHydrophobic",
                            "pos": None,
                            "len": 0,
                            "severity": sev,
                            "details": f"{max_consec_hydrophobic} consecutive hydrophobic",
                        }
                    )
                    total_score += sev
        results[seq] = {"score": total_score, "violations": violations}
    return results


def compute_liability_metrics(sequence, liability_modality, liability_peptide_type):
    metrics = {}
    # check if sequence is valid
    if not sequence or len(sequence) == 0:
        raise ValueError(f"Sequence is empty: '{sequence}'")

    liability_results = compute_liability_scores(
        [sequence],
        modality=liability_modality,
        peptide_type=liability_peptide_type,
    )
    liability_data = liability_results[sequence]

    # Store liability metrics
    metrics["liability_score"] = liability_data["score"]
    metrics["liability_num_violations"] = len(liability_data["violations"])

    # Count violations by severity
    high_severity_violations = [
        v for v in liability_data["violations"] if v["severity"] >= 10
    ]
    medium_severity_violations = [
        v for v in liability_data["violations"] if 5 <= v["severity"] < 10
    ]
    low_severity_violations = [
        v for v in liability_data["violations"] if v["severity"] < 5
    ]

    metrics["liability_high_severity_violations"] = len(high_severity_violations)
    metrics["liability_medium_severity_violations"] = len(medium_severity_violations)
    metrics["liability_low_severity_violations"] = len(low_severity_violations)

    # Count violations by type
    violation_counts = {}
    for v in liability_data["violations"]:
        motif = v["motif"]
        violation_counts[motif] = violation_counts.get(motif, 0) + 1

    # Store individual violation type counts as metrics
    for motif, count in violation_counts.items():
        metrics[f"liability_{motif}_count"] = count

    # Add detailed violation information
    # Group violations by type for intelligent reporting
    violations_by_type = {}
    for v in liability_data["violations"]:
        motif = v["motif"]
        if motif not in violations_by_type:
            violations_by_type[motif] = []
        violations_by_type[motif].append(v)

    # Initialize default values for all motifs to ensure consistent dataframe columns
    # Use the full motif set for the configured modality/peptide_type so columns are consistent
    all_motifs = set(
        get_motif_set(
            modality=liability_modality,
            peptide_type=liability_peptide_type,
        ).keys()
    )
    for motif in all_motifs:
        # Initialize all possible fields with default values
        metrics[f"liability_{motif}_count"] = 0
        metrics[
            f"liability_{motif}_position"
        ] = -1  # use -1 for no position (keeps int dtype)
        metrics[f"liability_{motif}_length"] = 0
        metrics[f"liability_{motif}_severity"] = 0
        metrics[f"liability_{motif}_details"] = ""
        metrics[f"liability_{motif}_positions"] = ""
        metrics[f"liability_{motif}_num_positions"] = 0
        metrics[f"liability_{motif}_global_details"] = ""
        metrics[f"liability_{motif}_avg_severity"] = 0.0

    # Store detailed violation information
    for motif, motif_violations in violations_by_type.items():
        if len(motif_violations) == 1:
            # Single violation - store all details
            v = motif_violations[0]
            # Ensure position is an integer; use -1 for non-positional violations
            metrics[f"liability_{motif}_position"] = (
                int(v["pos"]) if v["pos"] is not None else -1
            )
            metrics[f"liability_{motif}_length"] = v["len"]
            metrics[f"liability_{motif}_severity"] = v["severity"]
            if "details" in v:
                metrics[f"liability_{motif}_details"] = v["details"]
            else:
                metrics[f"liability_{motif}_details"] = ""
        else:
            # Multiple violations - store summary information
            positions = [v["pos"] for v in motif_violations if v["pos"] is not None]
            global_violations = [v for v in motif_violations if v["pos"] is None]

            if positions:
                # Store position range for positional violations
                metrics[f"liability_{motif}_positions"] = (
                    f"{min(positions)}-{max(positions)}"
                )
                metrics[f"liability_{motif}_num_positions"] = len(positions)
            else:
                metrics[f"liability_{motif}_positions"] = ""
                metrics[f"liability_{motif}_num_positions"] = 0

            if global_violations:
                # Store details for global violations
                details = [
                    v.get("details", "") for v in global_violations if v.get("details")
                ]
                if details:
                    metrics[f"liability_{motif}_global_details"] = "; ".join(details)
                else:
                    metrics[f"liability_{motif}_global_details"] = ""
            else:
                metrics[f"liability_{motif}_global_details"] = ""

            # Store average severity
            avg_severity = sum(v["severity"] for v in motif_violations) / len(
                motif_violations
            )
            metrics[f"liability_{motif}_avg_severity"] = round(avg_severity, 1)

    # Add a comprehensive violation summary for easy interpretation
    violation_summary = []
    for motif, motif_violations in violations_by_type.items():
        count = len(motif_violations)
        if count == 1:
            v = motif_violations[0]
            if v["pos"] is not None:
                violation_summary.append(f"{motif}(pos{v['pos']},sev{v['severity']})")
            else:
                details = v.get("details", "")
                violation_summary.append(
                    f"{motif}({details},sev{v['severity']})"
                    if details
                    else f"{motif}(sev{v['severity']})"
                )
        else:
            positions = [v["pos"] for v in motif_violations if v["pos"] is not None]
            if positions:
                violation_summary.append(
                    f"{motif}x{count}(pos{min(positions)}-{max(positions)},sev{motif_violations[0]['severity']})"
                )
            else:
                violation_summary.append(
                    f"{motif}x{count}(sev{motif_violations[0]['severity']})"
                )

    metrics["liability_violations_summary"] = (
        "; ".join(violation_summary) if violation_summary else ""
    )
    return metrics