posebench-fast 0.1.0__py3-none-any.whl

posebench_fast/metrics/aggregation.py

@@ -0,0 +1,388 @@
+"""Metrics aggregation and filtering for docking evaluation."""
+
+import copy
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+from posebench_fast.metrics.rmsd import (
+    TimeoutException,
+    get_symmetry_rmsd_with_isomorphisms,
+)
+
+
+def get_best_results_by_score(all_results, score_name):
+    """Get best sample for each UID based on a score.
+
+    Args:
+        all_results: Dictionary {uid: {'sample_metrics': [...]}}
+        score_name: Name of score to minimize (or 'random' for first sample)
+
+    Returns:
+        Dictionary {uid: best_sample_metrics}
+    """
+    filtered_results = {}
+
+    for uid in all_results:
+        metrics = all_results[uid]
+        if score_name == "random":
+            best_index = 0
+        else:
+            scores = np.array([metr[score_name] for metr in metrics["sample_metrics"]])
+            best_index = np.argmin(scores)
+
+        filtered_results[uid] = metrics["sample_metrics"][best_index]
+    return filtered_results
+
+
+def filter_results_by_posebusters(full_results, use_separate_samples=True):
+    """Filter samples to keep only those with maximum PoseBusters filters passed.
+
+    Args:
+        full_results: Dictionary of results
+        use_separate_samples: Whether samples are in 'sample_metrics' key
+
+    Returns:
+        Filtered results dictionary
+    """
+    for uid in full_results:
+        if use_separate_samples:
+            samples = full_results[uid]["sample_metrics"]
+        else:
+            samples = full_results[uid]
+
+        pb_filters_name = "posebusters_filters_passed_count"
+
+        scores = np.array([sample[pb_filters_name] for sample in samples])
+        best_score = max(scores)
+        filtered_samples = [
+            sample for sample in samples if sample[pb_filters_name] == best_score
+        ]
+        if use_separate_samples:
+            full_results[uid]["sample_metrics"] = filtered_samples
+        else:
+            full_results[uid] = filtered_samples
+    return full_results
+
+
+def filter_results_by_fast(full_results, use_separate_samples=True):
+    """Filter samples to keep only those with maximum fast PoseBusters filters passed.
+
+    Args:
+        full_results: Dictionary of results
+        use_separate_samples: Whether samples are in 'sample_metrics' key
+
+    Returns:
+        Filtered results dictionary
+    """
+    for uid in full_results:
+        if use_separate_samples:
+            samples = full_results[uid]["sample_metrics"]
+        else:
+            samples = full_results[uid]
+
+        try:
+            scores = np.array(
+                [sample["posebusters_filters_passed_count_fast"] for sample in samples]
+            )
+            best_score = max(scores)
+            filtered_samples = [
+                sample
+                for sample in samples
+                if sample["posebusters_filters_passed_count_fast"] == best_score
+            ]
+
+        except KeyError:
+            filtered_samples = samples
+
+        if use_separate_samples:
+            full_results[uid]["sample_metrics"] = filtered_samples
+        else:
+            full_results[uid] = filtered_samples
+    return full_results
+
+
+def filter_empty_results_and_keep_necessary_ids(
+    full_results, use_separate_samples=True, ids_to_keep=None
+):
+    """Filter out empty results and keep only specified IDs.
+
+    Args:
+        full_results: Dictionary of results
+        use_separate_samples: Whether samples are in 'sample_metrics' key
+        ids_to_keep: Optional list of UIDs to keep
+
+    Returns:
+        Filtered results dictionary
+    """
+    if ids_to_keep is not None:
+        all_pred_uids = {key.split("_mol")[0] for key in full_results}
+        uids_to_pop = [
+            f"{uid}_mol0" for uid in sorted(all_pred_uids - set(ids_to_keep))
+        ]
+    else:
+        uids_to_pop = []
+
+    if len(uids_to_pop) > 0:
+        print(f"Pop {len(uids_to_pop)} uids")
+
+    for uid in full_results:
+        if len(full_results[uid]) == 0:
+            print(f"{uid} has no valid samples")
+            uids_to_pop.append(uid)
+            continue
+
+        if use_separate_samples:
+            samples = full_results[uid]["sample_metrics"]
+        else:
+            samples = full_results[uid]
+
+        if len(samples) == 0:
+            print(f"{uid} has no valid samples")
+            uids_to_pop.append(uid)
+            continue
+
+    for uid in uids_to_pop:
+        full_results.pop(uid)
+
+    return full_results
+
+
+def get_final_results_for_df(
+    full_results,
+    score_names,
+    score_name_prefix="",
+    posebusters_filter=False,
+    fast_filter=False,
+    ids_to_keep=None,
+):
+    """Compute final metrics DataFrame from results.
+
+    Args:
+        full_results: Dictionary of results
+        score_names: List of score names to evaluate
+        score_name_prefix: Prefix for score names in output
+        posebusters_filter: Whether to also compute PoseBusters-filtered metrics
+        fast_filter: Whether to also compute fast-filter metrics
+        ids_to_keep: Optional list of UIDs to keep
+
+    Returns:
+        Tuple of (rows_list, all_scored_results)
+    """
+
+    def get_row(results, score_name, full_score_name, posebusters_filter):
+        scored_results = get_best_results_by_score(results, score_name)
+
+        rmsds = np.array([item["rmsd"] for item in scored_results.values()])
+        sym_rmsds = np.array([item["symm_rmsd"] for item in scored_results.values()])
+        tr_errs = np.array([item["tr_err"] for item in scored_results.values()])
+
+        row = {
+            "ranking": full_score_name,
+            "RMSD < 2A": (rmsds <= 2).mean(),
+            "RMSD < 5A": (rmsds <= 5).mean(),
+            "avg RMSD": rmsds.mean(),
+            "median RMSD": np.median(rmsds),
+            "SymRMSD < 2A": (sym_rmsds <= 2).mean(),
+            "SymRMSD < 5A": (sym_rmsds <= 5).mean(),
+            "avg SymRMSD": sym_rmsds.mean(),
+            "median SymRMSD": np.median(sym_rmsds),
+            "avg tr_err": tr_errs.mean(),
+            "median tr_err": np.median(tr_errs),
+            "tr_err < 1A": (tr_errs <= 1).mean(),
+            "num_samples": len(scored_results.values()),
+        }
+
+        if posebusters_filter:
+            posebusters_all = np.array(
+                [
+                    item["all_posebusters_filters_passed_count"]
+                    for item in scored_results.values()
+                ]
+            )
+            row["SymRMSD < 2A & PB valid"] = np.logical_and(
+                sym_rmsds < 2, posebusters_all == 27
+            ).mean()
+        return row, scored_results
+
+    rows_list = []
+    all_scored_results = {}
+
+    full_results = filter_empty_results_and_keep_necessary_ids(
+        full_results, use_separate_samples=True, ids_to_keep=ids_to_keep
+    )
+
+    if posebusters_filter:
+        filtered_results_posebusters = filter_results_by_posebusters(
+            copy.deepcopy(full_results)
+        )
+
+    if fast_filter:
+        filtered_results_fast = filter_results_by_fast(copy.deepcopy(full_results))
+
+    for score_name in score_names:
+        full_score_name = f"{score_name_prefix}{score_name}"
+
+        row, scored_results = get_row(
+            full_results,
+            score_name,
+            full_score_name,
+            posebusters_filter=posebusters_filter,
+        )
+        all_scored_results[full_score_name] = scored_results
+        rows_list.append(row)
+
+        if posebusters_filter:
+            real_score_name = f"{full_score_name}_posebusters"
+            row, scored_results = get_row(
+                filtered_results_posebusters,
+                score_name,
+                real_score_name,
+                posebusters_filter=posebusters_filter,
+            )
+            all_scored_results[real_score_name] = scored_results
+            rows_list.append(row)
+
+        if fast_filter:
+            real_score_name = f"{full_score_name}_fast"
+            row, scored_results = get_row(
+                filtered_results_fast,
+                score_name,
+                real_score_name,
+                posebusters_filter=posebusters_filter,
+            )
+            all_scored_results[real_score_name] = scored_results
+            rows_list.append(row)
+
+    return rows_list, all_scored_results
+
+
+def add_score_results(all_rmsds_new, score_res, score_name, n_samples=None):
+    """Add score results to samples.
+
+    Args:
+        all_rmsds_new: Dictionary of results
+        score_res: Score results dictionary
+        score_name: Type of score ('mult', 'bin', 'reg')
+        n_samples: Number of samples to use for averaging
+
+    Returns:
+        Extended results dictionary
+    """
+    extended_results = {}
+    for uid in tqdm(all_rmsds_new.keys(), desc="Adding score results"):
+        new_samples = []
+        for i in range(len(all_rmsds_new[uid])):
+            sample = all_rmsds_new[uid][i]
+            sample_scores = np.array(score_res[f"{uid}_{i}"])
+            nan_mask = np.isnan(sample_scores).sum(axis=1).astype(bool)
+            if nan_mask.sum() > 0:
+                if score_name == "mult":
+                    sample_scores[nan_mask, 2] = 6.0
+                    sample_scores[nan_mask, 0] = 0.0
+                    sample_scores[nan_mask, 1] = 0.0
+                elif score_name == "bin":
+                    sample_scores[nan_mask] = 0.0
+                elif score_name == "reg":
+                    sample_scores[nan_mask] = 50.0
+
+            sample_scores = -sample_scores
+            if n_samples is None:
+                n_samples = len(sample_scores)
+            mean_scores = np.mean(sample_scores[:n_samples], axis=0)
+
+            for idx in range(len(mean_scores)):
+                sample[f"{score_name}_{idx}"] = mean_scores[idx]
+
+            new_samples.append(sample)
+        extended_results[uid] = new_samples
+    return extended_results
+
+
+def get_simple_metrics_df(
+    all_real_rmsds, compute_symm_rmsd, mol2isomorphisms, score_names
+):
+    """Compute simple metrics DataFrame.
+
+    Args:
+        all_real_rmsds: Dictionary {uid: [samples]} with 'true_pos', 'transformed_orig'
+        compute_symm_rmsd: Whether to compute symmetry RMSD
+        mol2isomorphisms: Dictionary {uid: isomorphisms} from compute_all_isomorphisms
+        score_names: List of score names to include
+
+    Returns:
+        Tuple of (DataFrame, all_scored_results, full_results)
+    """
+    full_results = {}
+    for uid, samples in tqdm(all_real_rmsds.items(), desc="Computing metrics"):
+        samples_results = []
+        failed_symm_rmsd_count = 0
+
+        true_pos = samples[0]["true_pos"]
+        for idx in range(len(samples)):
+            pred_pos = samples[idx]["transformed_orig"]
+
+            if true_pos.shape[0] != pred_pos.shape[0]:
+                print(
+                    f"{uid}_{idx:<8} true_pos.shape[0] != pred_pos.shape[0]",
+                    true_pos.shape,
+                    pred_pos.shape,
+                )
+                continue
+
+            tr_pred = pred_pos.mean(axis=0)
+            tr_true = true_pos.mean(axis=0)
+            tr_err = np.linalg.norm(tr_pred - tr_true)
+
+            rmsd = np.sqrt(
+                ((true_pos - pred_pos) ** 2).sum(axis=1).sum() / true_pos.shape[0]
+            )
+            if compute_symm_rmsd and failed_symm_rmsd_count < 3:
+                try:
+                    mol2iso = mol2isomorphisms.get(uid.split("_conf")[0])
+                    if mol2iso is None:
+                        symm_rmsd = rmsd
+                        failed_symm_rmsd_count += 1
+                    else:
+                        symm_rmsd = get_symmetry_rmsd_with_isomorphisms(
+                            true_pos, pred_pos, mol2iso
+                        )
+                except TimeoutException:
+                    symm_rmsd = rmsd
+                    failed_symm_rmsd_count += 1
+            else:
+                symm_rmsd = rmsd
+
+            results = {
+                "tr_pred": tr_pred,
+                "tr_err": float(tr_err),
+                "symm_rmsd": float(symm_rmsd),
+                "rmsd": float(rmsd),
+                "pred_pos": pred_pos,
+            }
+            for score_name in set(score_names) - {"random", "symm_rmsd"}:
+                results[score_name] = float(samples[idx][score_name])
+
+            samples_results.append(results)
+        samples_results_dict = {
+            "sample_metrics": samples_results,
+            "true_pos": true_pos,
+            "orig_mol": samples[0]["orig_mol"],
+        }
+        if len(samples_results_dict["sample_metrics"]) > 0:
+            full_results[uid] = samples_results_dict
+        else:
+            print(f"{uid} has no valid samples")
+            print(
+                f"{uid} true_pos.shape[0] != pred_pos.shape[0]",
+                true_pos.shape,
+                pred_pos.shape,
+            )
+
+    if len(full_results) != len(all_real_rmsds):
+        print("Initial length of test_names", len(all_real_rmsds))
+        print("Length of full_results", len(full_results))
+
+    rows_list, all_scored_results = get_final_results_for_df(full_results, score_names)
+    return pd.DataFrame(rows_list), all_scored_results, full_results

posebench_fast/metrics/rmsd.py

@@ -0,0 +1,273 @@
+# Taken from https://github.com/RMeli/spyrmsd and https://github.com/gcorso/DiffDock/
+
+import signal
+from contextlib import contextmanager
+from typing import Any
+
+import numpy as np
+from spyrmsd import graph, molecule, qcp, utils
+
+
+class TimeoutException(Exception):  # noqa: N818
+    pass
+
+
+@contextmanager
+def time_limit(seconds):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timed out!")
+
+    signal.signal(signal.SIGALRM, signal_handler)
+    signal.alarm(seconds)
+    try:
+        yield
+    finally:
+        signal.alarm(0)
+
+
+def compute_all_isomorphisms(rdkit_mol):
+    """Compute all graph isomorphisms for a molecule.
+
+    Args:
+        rdkit_mol: RDKit molecule object
+
+    Returns:
+        List of isomorphism tuples (idx1, idx2)
+    """
+    try:
+        with time_limit(2):
+            mol = molecule.Molecule.from_rdkit(rdkit_mol)
+            G1 = graph.graph_from_adjacency_matrix(mol.adjacency_matrix, mol.atomicnums)
+            isomorphisms = graph.match_graphs(G1, G1)
+    except TimeoutException:
+        isomorphisms = [
+            (list(range(rdkit_mol.GetNumAtoms())), list(range(rdkit_mol.GetNumAtoms())))
+        ]
+    return isomorphisms
+
+
+def get_symmetry_rmsd_with_isomorphisms(coords1, coords2, isomorphisms):
+    """Compute symmetry-corrected RMSD using precomputed isomorphisms.
+
+    Args:
+        coords1: Reference coordinates (N, 3)
+        coords2: Query coordinates (N, 3)
+        isomorphisms: List of isomorphism tuples from compute_all_isomorphisms
+
+    Returns:
+        Minimum RMSD over all isomorphisms
+    """
+    with time_limit(1):
+        assert coords1.shape == coords2.shape
+
+        n = coords1.shape[0]
+        min_result = np.inf
+
+        for idx1, idx2 in isomorphisms:
+            c1i = coords1[idx1, :]
+            c2i = coords2[idx2, :]
+            result = np.sum((c1i - c2i) ** 2)
+
+            if result < min_result:
+                min_result = result
+
+        min_result = np.sqrt(min_result / n)
+        return min_result
+
+
+def get_symmetry_rmsd(mol, coords1, coords2, mol2=None, return_permutation=False):
+    """Compute symmetry-corrected RMSD.
+
+    Args:
+        mol: RDKit molecule
+        coords1: Reference coordinates
+        coords2: Query coordinates
+        mol2: Optional second molecule (if different from mol)
+        return_permutation: Whether to return the best permutation
+
+    Returns:
+        RMSD value (and permutation if requested)
+    """
+    with time_limit(10):
+        mol = molecule.Molecule.from_rdkit(mol)
+        mol2 = molecule.Molecule.from_rdkit(mol2) if mol2 is not None else mol2
+        mol2_atomicnums = mol2.atomicnums if mol2 is not None else mol.atomicnums
+        mol2_adjacency_matrix = (
+            mol2.adjacency_matrix if mol2 is not None else mol.adjacency_matrix
+        )
+        RMSD = symmrmsd(
+            coords1,
+            coords2,
+            mol.atomicnums,
+            mol2_atomicnums,
+            mol.adjacency_matrix,
+            mol2_adjacency_matrix,
+            return_permutation=return_permutation,
+        )
+        return RMSD
+
+
+def _rmsd_isomorphic_core(
+    coords1: np.ndarray,
+    coords2: np.ndarray,
+    aprops1: np.ndarray,
+    aprops2: np.ndarray,
+    am1: np.ndarray,
+    am2: np.ndarray,
+    center: bool = False,
+    minimize: bool = False,
+    isomorphisms: list[tuple[list[int], list[int]]] | None = None,
+    atol: float = 1e-9,
+) -> tuple[float, list[tuple[list[int], list[int]]], tuple[list[int], list[int]]]:
+    """
+    Compute RMSD using graph isomorphism.
+
+    Parameters
+    ----------
+    coords1: np.ndarray
+        Coordinate of molecule 1
+    coords2: np.ndarray
+        Coordinates of molecule 2
+    aprops1: np.ndarray
+        Atomic properties for molecule 1
+    aprops2: np.ndarray
+        Atomic properties for molecule 2
+    am1: np.ndarray
+        Adjacency matrix for molecule 1
+    am2: np.ndarray
+        Adjacency matrix for molecule 2
+    center: bool
+        Centering flag
+    minimize: bool
+        Compute minized RMSD
+    isomorphisms: Optional[List[Dict[int,int]]]
+        Previously computed graph isomorphism
+    atol: float
+        Absolute tolerance parameter for QCP (see :func:`qcp_rmsd`)
+
+    Returns
+    -------
+    Tuple[float, List[Dict[int, int]]]
+        RMSD (after graph matching) and graph isomorphisms
+    """
+
+    assert coords1.shape == coords2.shape
+
+    n = coords1.shape[0]
+
+    c1 = utils.center(coords1) if center or minimize else coords1
+    c2 = utils.center(coords2) if center or minimize else coords2
+
+    if isomorphisms is None:
+        G1 = graph.graph_from_adjacency_matrix(am1, aprops1)
+        G2 = graph.graph_from_adjacency_matrix(am2, aprops2)
+        isomorphisms = graph.match_graphs(G1, G2)
+
+    min_result = np.inf
+    min_isomorphisms = None
+
+    for idx1, idx2 in isomorphisms:
+        c1i = c1[idx1, :]
+        c2i = c2[idx2, :]
+
+        if not minimize:
+            result = np.sum((c1i - c2i) ** 2)
+        else:
+            result = qcp.qcp_rmsd(c1i, c2i, atol)
+
+        if result < min_result:
+            min_result = result
+            min_isomorphisms = (idx1, idx2)
+
+    if not minimize:
+        min_result = np.sqrt(min_result / n)
+
+    return min_result, isomorphisms, min_isomorphisms
+
+
+def symmrmsd(
+    coordsref: np.ndarray,
+    coords: np.ndarray | list[np.ndarray],
+    apropsref: np.ndarray,
+    aprops: np.ndarray,
+    amref: np.ndarray,
+    am: np.ndarray,
+    center: bool = False,
+    minimize: bool = False,
+    cache: bool = True,
+    atol: float = 1e-9,
+    return_permutation: bool = False,
+) -> Any:
+    """
+    Compute RMSD using graph isomorphism for multiple coordinates.
+
+    Parameters
+    ----------
+    coordsref: np.ndarray
+        Coordinate of reference molecule
+    coords: List[np.ndarray]
+        Coordinates of other molecule
+    apropsref: np.ndarray
+        Atomic properties for reference
+    aprops: np.ndarray
+        Atomic properties for other molecule
+    amref: np.ndarray
+        Adjacency matrix for reference molecule
+    am: np.ndarray
+        Adjacency matrix for other molecule
+    center: bool
+        Centering flag
+    minimize: bool
+        Minimum RMSD
+    cache: bool
+        Cache graph isomorphisms
+    atol: float
+        Absolute tolerance parameter for QCP (see :func:`qcp_rmsd`)
+
+    Returns
+    -------
+    float: Union[float, List[float]]
+        Symmetry-corrected RMSD(s) and graph isomorphisms
+    """
+
+    if isinstance(coords, list):
+        RMSD: Any = []
+        isomorphism = None
+        min_iso = []
+
+        for c in coords:
+            if not cache:
+                isomorphism = None
+
+            srmsd, isomorphism, min_i = _rmsd_isomorphic_core(
+                coordsref,
+                c,
+                apropsref,
+                aprops,
+                amref,
+                am,
+                center=center,
+                minimize=minimize,
+                isomorphisms=isomorphism,
+                atol=atol,
+            )
+            min_iso.append(min_i)
+            RMSD.append(srmsd)
+
+    else:
+        RMSD, isomorphism, min_iso = _rmsd_isomorphic_core(
+            coordsref,
+            coords,
+            apropsref,
+            aprops,
+            amref,
+            am,
+            center=center,
+            minimize=minimize,
+            isomorphisms=None,
+            atol=atol,
+        )
+
+    if return_permutation:
+        return RMSD, min_iso
+    return RMSD

posebench_fast/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility functions for posebench-fast."""