from posecheck.posecheck import PoseCheck

posecheck/__init__.py
cch1999/posecheck
import glob
import os
from typing import List, Tuple

import datamol as dm
import pandas as pd
import prolif as plf
import torch
from tqdm import tqdm

from posecheck.utils.clashes import count_clashes_list
from posecheck.utils.interactions import generate_interaction_df
from posecheck.utils.loading import (get_ids_to_pockets, get_pdbqt_mol,
                               load_mols_from_sdf, load_protein_from_pdb,
                               read_pdbqt)
from posecheck.utils.strain import get_strain_energy

# get directory of this file
DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(DIR, "../../data/")

POCKET_DIR = os.path.join(DATA_DIR, "crossdocked/")
PDBQT_DIR = os.path.join(DATA_DIR, "crossdocked_pdbqt/")


class BasePosesDataset(object):
    """Base class for loading protein-ligand complexes for PoseCheck."""

    def __init__(
        self,
        data_dir: str = DATA_DIR,
        split_path: str = os.path.join(DATA_DIR, "split_by_name.pt"),
        pocket_dir: str = POCKET_DIR,
        pdbqt_dir: str = PDBQT_DIR,
        protein_data=None,
    ):
        """
        Initializes the BasePosesDataset class.

        Args:
            data_dir (str, optional): The directory where the data is located. Defaults to DATA_DIR.
            split_path (str, optional): The path to the file that contains the train, validation, and test splits. Defaults to os.path.join(DATA_DIR, "split_by_name.pt").
            pocket_dir (str, optional): The directory where the protein pocket data is located. Defaults to POCKET_DIR.
            pdbqt_dir (str, optional): The directory where the PDBQT files are located. Defaults to PDBQT_DIR.
            protein_data (optional): Preloaded protein data. If None, the protein data will be loaded from the proteins.pt file in the data directory. Defaults to None.
        """
        self.data_dir = data_dir
        self.pocket_dir = pocket_dir
        self.pdbqt_dir = pdbqt_dir
        split_path = os.path.join(data_dir, "crossdocked_split_by_name.pkl")
        self.names = get_ids_to_pockets(split_path)

        if protein_data is None:
            proteins_path = os.path.join(data_dir, "proteins.pt")
            if os.path.exists(proteins_path):
                self.proteins = torch.load(proteins_path)
            else:
                raise Exception("No proteins found, preprocess first")

        else:
            self.proteins = protein_data

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}"

    def __len__(self):
        return len(self.names)

    def get_pdb_path(self, idx: int):
        return glob.glob(
            self.pocket_dir + f"{self.names[idx]}*" + "*.pdb", recursive=True
        )[0]

    def get_pdbqt_path(self, idx: int):
        return glob.glob(
            self.pdbqt_dir + f"{self.names[idx]}*" + "*.pdbqt", recursive=True
        )[0]

    def prepare_pdb(self, idx: int):
        # pdb_path = glob.glob(self.pocket_dir + f"{name}*" + "*.pdb", recursive=True)[0]
        # return load_protein_from_pdb(pdb_path)
        return self.proteins[idx]

    def prepare_sdf(self, name: str, idx: int = 0):
        """Implement in child class. This should be the only class that needs defining for a new dataset."""
        raise NotImplementedError

    def __getitem__(self, idx) -> Tuple[plf.Molecule, plf.Molecule]:
        """Returns a tuple of protein and ligand structures.

        Args:
            idx (int): The index of the pocket.

        Returns:
            Tuple[pd.AtomGroup, pybel.Molecule]: The protein and ligand structures.
        """

        # Select pocket
        name = self.names[idx]

        # Load protein and ligand
        prot = self.prepare_pdb(idx)
        lig = self.prepare_sdf(name, idx)

        return prot, lig

    def get_interaction_df(self, idx: int) -> pd.DataFrame:
        """Returns a dataframe of interactions between the protein and ligand.

        Args:
            idx (int): The index of the pocket.

        Returns:
            pd.DataFrame: The interactions between the protein and ligand.
        """
        try:
            prot, lig = self.__getitem__(idx)
            interaction_df = generate_interaction_df(prot, lig)
            return interaction_df
        except Exception as e:
            print(f"Error with {idx}: {str(e)}")
            return None

    @property
    def all_interactions(self) -> List[pd.DataFrame]:
        """Returns a list of interaction DataFrames for all pockets in this dataset.

        Calculates the interactions if they have not already been calculated.

        List will contain None values if there was an error with the pocket.

        Returns:
            List[pd.DataFrame]: A list of interaction DataFrames.
        """
        interactions_path = (
            self.data_dir
            + f"/benchmarks/preprocessed/{self.__repr__()}_interactions.pt"
        )

        if not os.path.exists(interactions_path):
            print("Calculating interactions...")
            dataset_name = self.__repr__()
            interactions = [
                self.get_interaction_df(idx)
                for idx in tqdm(range(len(self)), desc=dataset_name)
            ]
            torch.save(interactions, interactions_path)
        else:
            interactions = torch.load(interactions_path)

        # Count number of None values
        n_nones = sum([1 for i in interactions if i is None])
        print(f"{self.__repr__()}: Number of None values: {n_nones}")

        return interactions

    @property
    def all_clashes(self):
        clashes_path = (
            self.data_dir + f"/benchmarks/clashes/{self.__repr__()}_clashes.pt"
        )

        if not os.path.exists(clashes_path):
            print("Calculating clashes...")
            dataset_name = self.__repr__()
            clashes = [
                count_clashes_list(*self.__getitem__(idx), target=self.names[idx][:4])
                for idx in tqdm(range(len(self)), desc=dataset_name)
            ]
            # clashes = torch.stack(torch.tensor(clashes))
            torch.save(clashes, clashes_path)
        else:
            clashes = torch.load(clashes_path)

        # Count number of None values
        # n_nones = sum([1 for i in interactions if i is None])
        # print(f"{self.__repr__()}: Number of None values: {n_nones}")

        return clashes

    @property
    def all_strain_energy(self):
        energies_path = (
            self.data_dir + f"/benchmarks/energies/{self.__repr__()}_energies.pt"
        )

        if not os.path.exists(energies_path):
            mols = []
            for idx in tqdm(range(len(self))):
                mols.extend(self.load_mols(idx))

            energies = [get_strain_energy(mol) for mol in tqdm(mols)]

            torch.save(energies, energies_path)
        else:
            energies = torch.load(energies_path)

        return energies


# class CrossdockedDataset(BasePosesDataset):
#     def __init__(
#         self,
#         data_dir: str = DATA_DIR
#         split_path: str = "/Users/charlie/projects/poses_benchmark/data/split_by_name.pt",
#     ):
#         super().__init__(data_dir=data_dir, split_path=split_path)
#         self.ligand_dir = data_dir + "/crossdocked/"

#     def prepare_sdf(self, name: str, idx: int = 0):
#         sdf_path = glob.glob(self.ligand_dir + f"{name}*" + "*.sdf", recursive=True)[0]
#         return 1(sdf_path)


# class ZINCSamples(BasePosesDataset):
#     def __init__(
#         self,
#         data_dir: str = "/Users/charlie/projects/poses_benchmark/data",
#         split_path: str = "/Users/charlie/projects/poses_benchmark/data/split_by_name.pt",
#     ):
#         super().__init__(data_dir=data_dir, split_path=split_path)
#         self.ligand_dir = data_dir + "/benchmarks/zinc/"

#     def prepare_sdf(self, name: str, idx: int = 0):
#         sdf_path = glob.glob(self.ligand_dir + f"{name}*" + "*.sdf", recursive=True)[0]
#         return load_mols_from_sdf(sdf_path)


# class TargetDiffSamples(BasePosesDataset):
#     def __init__(
#         self,
#         data_dir: str = "/Users/charlie/projects/poses_benchmark/data",
#         split_path: str = "/Users/charlie/projects/poses_benchmark/data/split_by_name.pt",
#     ):
#         super().__init__(data_dir=data_dir, split_path=split_path)

#         self.ligand_dir = data_dir + "/benchmarks/targetdiff/"
#         self.data = torch.load(
#             os.path.join(self.ligand_dir, "targetdiff_vina_docked.pt")
#         )

#     def prepare_sdf(self, name: str, idx: int = 0):
#         samples = self.data[idx]

#         mols = [sample["mol"] for sample in samples]
#         tmp_path = self.ligand_dir + "tmp.sdf"
#         dm.to_sdf(mols, tmp_path)

#         lig = load_mols_from_sdf(tmp_path)
#         os.remove(tmp_path)

#         return lig


# class BenchmarkModelSamples(BasePosesDataset):
#     def __init__(
#         self,
#         model_name: str,  # cvae, p2m or sbdd
#         data_dir: str = DATA_DIR,
#         split_path: str = os.path.join(DATA_DIR, "split_by_name.pt"),
#     ):
#         super().__init__(data_dir=data_dir, split_path=split_path)

#         self.name = model_name
#         self.ligand_dir = data_dir + f"/benchmarks/{model_name}/"

#     def __repr__(self) -> str:
#         return f"{self.name}"

#     def prepare_sdf(self, name: str, idx: int = 0):
#         sample_dir = self.ligand_dir + f"pocket_{idx}/"

#         # Remove tmp.sdf incase it exists
#         tmp_path = sample_dir + "tmp.sdf"
#         if os.path.exists(tmp_path):
#             os.remove(tmp_path)

#         # Read all SDFs in the folder and make a single temp SDF with all molecules
#         sdf_paths = glob.glob(sample_dir + "*.sdf", recursive=True)

#         # If no SDFs found, return None
#         if len(sdf_paths) == 0:
#             print("No SDFs found for", name, idx, "returning None")
#             return None

#         # Read all SDFs and make a single temp SDF with all molecules
#         mols = []
#         for sdf_path in sdf_paths:
#             try:
#                 mols.append(dm.read_sdf(sdf_path)[0])
#             except:
#                 pass

#         dm.to_sdf(mols, tmp_path)

#         # Load the temp SDF
#         ligs = load_mols_from_sdf(tmp_path)
#         os.remove(tmp_path)

#         return ligs


class DockedMolsDataset(BasePosesDataset):
    def __init__(
        self,
        name: str,
        docked: bool,
        data_dir: str = DATA_DIR,
        split_path: str = os.path.join(DATA_DIR, "split_by_name.pt"),
    ):
        """Dataset for loading docked molecules from dataset:

        Args:
            name (str): Name of dataset
            docked (bool): Whether to load docked or generated molecules
            data_dir (str, optional): [description]. Defaults to DATA_DIR.
            split_path (str, optional): [description]. Defaults to os.path.join(DATA_DIR, "split_by_name.pt").
        """

        super().__init__(data_dir=data_dir, split_path=split_path)
        self.name = name
        self.docked = docked

        self.ligand_dir = (
            data_dir + "/benchmarks/targetdiff/"
        )  # NOTE samples taken from targetdiff

        #
        self.data = torch.load(os.path.join(self.ligand_dir, f"{name}_vina_docked.pt"))

    def __repr__(self) -> str:
        return f"{self.name}_{'docked' if self.docked else 'generated'}"

    def load_generated_mols(self, idx: int = 0):
        if self.name != "crossdocked_test":
            return [sample["mol"] for sample in self.data[idx]]
        else:
            return [self.data[idx]["mol"]]

    def load_docked_mols(self, idx: int = 0):
        if self.name != "crossdocked_test":
            return [
                get_pdbqt_mol(sample["vina"]["dock"][0]["pose"])
                for sample in self.data[idx]
            ]
        else:
            return [
                get_pdbqt_mol(self.data[idx]["vina"]["dock"][0]["pose"])
            ]  # Because crossdocked_test only has 1 sample

    def load_mols(self, idx: int):
        """Loads molecules for a target directly without computing interactions or preparing ligand with plip.

        Args:
            idx (int, optional): id. Defaults to 0.

        Returns:
            List[Chem.Mol]: List of molecules.
        """
        if self.docked:
            return self.load_docked_mols(idx)
        elif not self.docked:
            return self.load_generated_mols(idx)

    def prepare_sdf(self, name: str, idx: int = 0):
        mols = self.load_mols(idx)

        # Remove None values
        mols = [dm.add_hs(mol, add_coords=True) for mol in mols if mol is not None]

        tmp_path = self.ligand_dir + "tmp.sdf"
        dm.to_sdf(mols, tmp_path)

        try:
            lig = load_mols_from_sdf(tmp_path)
            os.remove(tmp_path)
            return lig
        except Exception as e:
            print(e)
            return []


class DiffSBDDSamples(BasePosesDataset):
    def __init__(
        self,
        docked,
        diffsbdd_dir="/home/cch57/projects/poses_benchmark/data/benchmarks/diffsbdd/",
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.name = "diffsbdd"
        self.docked = docked
        self.ligand_dir = diffsbdd_dir

    def __repr__(self) -> str:
        return f"{self.name}_{'docked' if self.docked else 'generated'}"

    def load_generated_mols(self, idx: int = 0):
        name = self.names[idx]
        pdbqt_files = glob.glob(self.ligand_dir + f"*{name}*" + "*.pdbqt")
        pdbqt_files = [
            pdbqt_file for pdbqt_file in pdbqt_files if "out" not in pdbqt_file
        ]
        generated_mols = [read_pdbqt(pdbqt_file) for pdbqt_file in pdbqt_files]
        return generated_mols

    def load_docked_mols(self, idx: int = 0):
        name = self.names[idx]
        pdbqt_files = glob.glob(self.ligand_dir + f"*{name}*" + "*.pdbqt")
        pdbqt_files = [
            pdbqt_file for pdbqt_file in pdbqt_files if "out" not in pdbqt_file
        ]

        docked_paths = [
            pdbqt_file.split(".")[0] + "_out.pdbqt" for pdbqt_file in pdbqt_files
        ]
        docked_mols = [read_pdbqt(pdbqt_file) for pdbqt_file in docked_paths]
        return docked_mols

    def load_mols(self, idx: int = 0):
        if self.docked:
            return self.load_docked_mols(idx)
        elif not self.docked:
            return self.load_generated_mols(idx)

    def prepare_sdf(self, name: str, idx: int = 0):
        mols = self.load_mols(idx)

        # Remove None values
        mols = [dm.add_hs(mol, add_coords=True) for mol in mols if mol is not None]

        tmp_path = self.ligand_dir + "tmp.sdf"
        dm.to_sdf(mols, tmp_path)

        try:
            lig = load_mols_from_sdf(tmp_path)
            os.remove(tmp_path)
            return lig
        except Exception as e:
            print(e)
            return []


class GraphBPSamples(BasePosesDataset):
    def __init__(self, data_dir: str = DATA_DIR, *args, **kwargs):
        super().__init__(data_dir=data_dir, *args, **kwargs)
        self.name = "graphbp"
        self.data_dir = os.path.join(data_dir, "benchmarks", self.name)

        self.sdf_files = [
            glob.glob(
                os.path.join(self.data_dir, f"*{name[:4].replace('-', '_')}*.sdf")
            )
            for name in self.names
        ]

    def load_mols(self, idx: int = 0):
        target_sdfs = self.sdf_files[idx]

        mols = []
        for sdf in target_sdfs:
            try:
                mols.append(dm.read_sdf(sdf)[0])
            except:
                print(f"Failed to load {sdf} adding None")
                mols.append(None)

        # ! NOTE in the middle of refactoring to remove self.docked
        # if self.docked:
        #     return self.load_docked_mols(idx)
        # elif not self.docked:
        #     return self.load_generated_mols(idx)

        return mols

    def prepare_sdf(self, name: str, idx: int = 0):
        raise NotImplementedError


class FLAGSamples(BasePosesDataset):
    def __init__(self, data_dir: str = DATA_DIR, *args, **kwargs):
        super().__init__(data_dir=data_dir, *args, **kwargs)
        self.name = "flag"
        # self.data_dir = os.path.join(data_dir, 'benchmarks', self.name)
        self.data_dir
        self.ligand_dir = data_dir + "/benchmarks/flag/"
        self.docked = False

        self.sdf_files = [
            glob.glob(
                os.path.join(self.ligand_dir, f"sample_{i}", "generated_mols.sdf")
            )
            for i in range(len(self.names))
        ]

    def load_mols(self, idx: int = 0):
        # TODO Make for generated too
        mols = dm.read_sdf(self.sdf_files[idx][0])

        return mols

    def prepare_sdf(self, name: str, idx: int = 0):
        mols = self.load_mols(idx)

        mols = [dm.add_hs(mol, add_coords=True) for mol in mols if mol is not None]

        # Remove None values
        # mols = [dm.add_hs(mol, add_coords=True) for mol in mols if mol is not None]

        tmp_path = self.ligand_dir + "tmp.sdf"
        dm.to_sdf(mols, tmp_path)

        try:
            lig = load_mols_from_sdf(tmp_path)
            os.remove(tmp_path)
            return lig
        except Exception as e:
            print(e)
            return []


DECOMPDIFF_DIR = "/home/cch57/rds/rds-liolab-8eFRJ5nImu8/data/crossdocked"


class DecompDiffSamples(BasePosesDataset):
    def __init__(self, data_dir: str = DECOMPDIFF_DIR, *args, **kwargs):
        super().__init__(data_dir=data_dir, *args, **kwargs)
        self.name = "decompdiff"
        # self.data_dir = os.path.join(data_dir, 'benchmarks', self.name)
        self.data_dir  # = data_dir
        self.ligand_dir = data_dir + "/benchmarks/flag/"
        self.docked = False

        # self.sdf_files = [glob.glob(os.path.join(self.ligand_dir, f'sample_{i}', "generated_mols.sdf")) for i in range(len(self.names))]

    def load_mols(self, idx: int = 0):
        # get all directories with 'sampling_drift' in the name
        name = self.names[idx][:4]

        paths = glob.glob(
            os.path.join(self.data_dir, "sampling_drift" + f"*{name}*", "*.pt")
        )

        if len(paths) == 0:
            print(f"No paths found for {name}")
            return []
        else:
            mols_path = glob.glob(
                os.path.join(self.data_dir, "sampling_drift" + f"*{name}*", "*.pt")
            )[0]

        data = torch.load(mols_path)

        mols = [sample["mol"] for sample in data]
        mols

        return mols

    def prepare_sdf(self, name: str, idx: int = 0):
        mols = self.load_mols(idx)

        mols = [dm.add_hs(mol, add_coords=True) for mol in mols if mol is not None]

        # Remove None values
        # mols = [dm.add_hs(mol, add_coords=True) for mol in mols if mol is not None]

        tmp_path = self.ligand_dir + "tmp.sdf"
        dm.to_sdf(mols, tmp_path)

        try:
            lig = load_mols_from_sdf(tmp_path)
            os.remove(tmp_path)
            return lig
        except Exception as e:
            print(e)
            return []


def load_docking_data():
    print("Loading docking data...")

    DOCKING_DIR = "/home/cch57/projects/poses_benchmark/data/docking_results"

    models = os.listdir(DOCKING_DIR)
    dfs = []

    for model in models:
        docking_files = glob.glob(os.path.join(DOCKING_DIR, model, "*.pt"))
        data = [torch.load(f) for f in docking_files]

        for i, target in enumerate(data):
            for j, score in enumerate(target["raw"]):
                try:
                    if len(data[i]["raw"][j]["redocked"]["mols"]) == 0:
                        data[i]["raw"][j]["redocked"]["mols"] = [None] * 8
                except:
                    print("hit None - skipping")

        df = pd.DataFrame(
            {
                "generated": [
                    score["raw"][index]["score_only"]["affinity"]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                "minimized": [
                    score["raw"][index]["minimized"]["min_affinity"]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                "redocked_best": [
                    score["raw"][index]["redocked"]["affinities"][0]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                "redocked_all": [
                    score["raw"][index]["redocked"]["affinities"]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                # mols
                "generated_mol": [
                    score["raw"][index]["score_only"]["mol"]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                "minimized_mol": [
                    score["raw"][index]["minimized"]["mol"]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                "redocked_best_mol": [
                    score["raw"][index]["redocked"]["mols"][0]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                #'redocked_all_mol' : [score['raw'][index]['redocked']['mols'] for score in data for index in range(len(score['raw'])) if score['raw'][index]['score_only'] is not None],
                # rmsds
                "min_rmsd": [
                    score["raw"][index]["minimized"]["min_rmsd"]
                    for score in data
                    for index in range(len(score["raw"]))
                    if score["raw"][index]["score_only"] is not None
                ],
                "target": [
                    score["receptor_path"].split("/")[-1].split(".")[0][:4]
                    for score in data
                    for _ in range(len(score["raw"]))
                    if score["raw"][_]["score_only"] is not None
                ],
            }
        )

        df["method"] = model

        dfs.append(df)

    df = pd.concat(dfs)

    print("Done")
    return df


# DOCKING_DATA = load_docking_data()
DOCKING_DATA = torch.load(os.path.join(DATA_DIR, "docking_data.pt"))
PROTEIN_DATA = torch.load(os.path.join(DATA_DIR, "proteins.pt"))


class RedockedData(BasePosesDataset):
    def __init__(
        self, name: str,
                type: str,
                data_dir: str = DATA_DIR,
                pdbqt_dir: str = PDBQT_DIR,
                docking_data=DOCKING_DATA,
                protein_data=PROTEIN_DATA
    ):
        super().__init__(protein_data=protein_data, data_dir=data_dir, pdbqt_dir=pdbqt_dir)

        self.name = name
        self.type = type
        self.data = docking_data[docking_data["method"] == name]
        assert len(self.data) > 0, "No data found"

    def __len__(self):
        return len(self.names)

    def __repr__(self) -> str:
        return f"{self.name}_{self.type}"

    def load_mols(self, idx: int):
        name = self.names[idx][:4]

        if self.type == "score_only":
            mols = self.data[self.data["target"] == name]["generated_mol"].to_list()
        elif self.type == "minimized":
            mols = self.data[self.data["target"] == name]["minimized_mol"].to_list()
        elif self.type == "redocked":
            mols = self.data[self.data["target"] == name]["redocked_best_mol"].to_list()
        else:
            raise ValueError("Invalid type")

        return [mol for mol in mols if mol is not None]

    def prepare_sdf(self, name: str, idx: int = 0):
        mols = self.load_mols(idx)

        # Remove None values
        mols = [dm.add_hs(mol, add_coords=True) for mol in mols if mol is not None]

        tmp_path = f"tmp_ligands_{self.name}.sdf"
        dm.to_sdf(mols, tmp_path)

        try:
            lig = load_mols_from_sdf(tmp_path)
            os.remove(tmp_path)
            return lig
        except Exception as e:
            print(e)
            os.remove(tmp_path)
            return []

posecheck/data/datasets.py
cch1999/posecheck
import argparse
import os

from rdkit import Chem
from tqdm import tqdm

import wandb
from posecheck.data.datasets import *
from posecheck.utils.docking import SMINA
from posecheck.utils.docking_utils import *


def benchmark_mols_for_target(
    mols: List[Chem.Mol], receptor_path: str
) -> Dict[str, Dict[str, float]]:
    """Benchmark a list of molecules for a target.

    Parameters
    ----------
    mols : list of rdkit.Chem.Mol
        The molecules to benchmark.
    target : str
        The target to benchmark against.

    Returns
    -------
    dict of str: dict of str: float
        A dictionary containing the results of the benchmark.
    """

    # Define the docking software and set the receptor
    smina = SMINA()
    smina.set_receptor(receptor_path)

    # Calculate the results for the raw molecules
    print("Calculating results for raw molecules")
    raw_results = [smina.calculate_all(mol) for mol in tqdm(mols) if mol is not None]

    # # Calculate the results for the relaxed molecules
    # print('Calculating results for relaxed molecules')
    # relaxed_results = []
    # for mol in mols:
    #     try:
    #         relaxed_mol = relax_mol(mol)
    #         relaxed_results.append(smina.calculate_all(relaxed_mol))
    #     except:
    #         relaxed_results.append(None)

    relaxed_results = [None] * len(raw_results)

    # relaxed_mols = [relax_mol(mol) for mol in mols]
    # relaxed_results = [smina.calculate_all(mol) for mol in tqdm(relaxed_mols)]

    results = {
        "raw": raw_results,
        "relaxed": relaxed_results,
        "receptor_path": receptor_path,
    }

    # Return the results
    return results


if __name__ == "__main__":
    args = argparse.ArgumentParser()
    args.add_argument("--dataset_name", type=str, default="decompdiff")
    args.add_argument(
        "--data_dir", type=str, default="/home/cch57/projects/poses_benchmark/data"
    )
    args = args.parse_args()

    data_dir = args.data_dir
    dataset_name = args.dataset_name

    # Load the dataset
    if dataset_name == "diffsbdd":
        dataset = DiffSBDDSamples(docked=False)
    elif dataset_name == "flag":
        dataset = FLAGSamples()
    elif dataset_name == "decompdiff":
        dataset = DecompDiffSamples()
    else:
        dataset = DockedMolsDataset(dataset_name, docked=False)

    # Set up the output directory and logger
    out_dir = os.path.join(data_dir, "docking_results", dataset_name)
    logger = wandb.init(project="poses_benchmark", entity="cch1999", name=dataset_name)

    # Make the output directory if it doesn't exist
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    # Loop over the dataset and calculate the docking results
    for idx in range(len(dataset)):
        logger.log({"target": idx})
        print(f"Calculating docking results for target {idx}")

        out_path = os.path.join(out_dir, f"target_{idx}.pt")
        if os.path.exists(out_path):  # Skip if already done
            continue

        receptor_path = dataset.get_pdbqt_path(idx)
        mols = dataset.load_mols(idx)

        print(f"Loaded {len(mols)} molecules ****")

        # Calculate the docking results
        docking_results = benchmark_mols_for_target(mols, receptor_path)

        # Save the results
        torch.save(docking_results, out_path)

posecheck/docking_eval.py
cch1999/posecheck
import argparse

from posecheck.data.datasets import *

if __name__ == "__main__":
    args = argparse.ArgumentParser()
    args.add_argument("--dataset", type=str, default="crossdocked_test")
    args = args.parse_args()

    dataset_name = args.dataset

    # Load the dataset
    types = ["score_only", "minimized", "redocked"]

    for method in [
        "ar",
        "targetdiff",
        "pocket2mol",
        "cvae",
        "diffsbdd",
        "flag",
        "decompdiff",
        "pocket2mol",
    ]:
        for type in types:
            dataset = RedockedData(method, type)
            try:
                dataset.all_strain_energy
            except:
                pass

            try:
                dataset.all_interactions
            except:
                pass

            try:
                dataset.all_clashes
            except:
                pass

posecheck/eval.py
cch1999/posecheck
import os
from typing import List

import datamol as dm
import pandas as pd
from rdkit import Chem

from posecheck.utils.chem import rmsd
from posecheck.utils.clashes import count_clashes
from posecheck.utils.interactions import generate_interaction_df
from posecheck.utils.loading import (load_mols_from_rdkit, load_mols_from_sdf,
                               load_protein_from_pdb, read_pdbqt)
from posecheck.utils.strain import get_strain_energy


class PoseCheck(object):
    """Main class for PoseCheck.

    Load a protein-ligand complex, check for clashes, strain, interactions, etc.

    Example usage:
    >>> pc = PoseCheck()
    >>> pc.load_protein_from_pdb("data/1a2b.pdb")
    >>> pc.load_ligand_from_sdf("data/1a2b_ligand.sdf")
    # Run posecheck
    >>> pc.check_clashes()
    >>> pc.check_strain()
    >>> pc.check_interactions()
    # Or run all at once
    >>> pc.run()
    """

    def __init__(
        self, reduce_path: str = "reduce", clash_tolerance: float = 0.5
    ) -> None:
        """Initialize the PoseCheck class.

        Args:
            reduce_path (str, optional): The path to the reduce executable. Defaults to "reduce".
            clash_tolerance (float, optional): The clash tolerance for checking clashes. Defaults to 0.5 A.
        """
        self.reduce_path = reduce_path
        self.clash_tolerance = clash_tolerance

    def load_protein_from_pdb(self, pdb_path: str) -> None:
        """Load a protein from a PDB file.

        Args:
            pdb_path (str): The path to the PDB file.

        Returns:
            None
        """
        self.protein = load_protein_from_pdb(pdb_path, reduce_path=self.reduce_path)

    def load_ligands_from_sdf(self, sdf_path: str) -> None:
        """Load a ligand from an SDF file."""
        self.ligands = load_mols_from_sdf(sdf_path)

    def load_ligands_from_pdbqt(self, pdbqt_path: str) -> None:
        """Load a ligand from a PDBQT file."""
        mol = read_pdbqt(pdbqt_path)

        # Save to tmp sdf file and load with Hs
        tmp_path = pdbqt_path.split(".pdbqt")[0] + "_tmp.pdb"
        dm.to_sdfile(mol, tmp_path)
        self.ligands = load_mols_from_sdf(tmp_path)
        os.remove(tmp_path)

    def load_ligands_from_mols(self, mols: List[Chem.Mol], add_hs: bool = True) -> None:
        """Load ligands from a list of RDKit mols.

        Args:
            mols (List[Chem.Mol]): The list of RDKit mol objects representing the ligands.
            add_hs (bool, optional): Whether to add hydrogens to the ligands. Defaults to True.

        Returns:
            None
        """
        self.ligands = load_mols_from_rdkit(mols, add_hs=add_hs)

    def load_ligands(self, ligand) -> None:
        """Detect ligand type and load.

        Args:
            ligand (str or Chem.Mol): The ligand to load.

        Raises:
            ValueError: If the ligand type is unknown.

        Returns:
            None
        """

        if isinstance(ligand, str):
            if ligand.endswith(".sdf"):
                self.load_ligand_from_sdf(ligand)
            elif ligand.endswith(".pdbqt"):
                self.load_ligand_from_pdbqt(ligand)
            else:
                raise ValueError("Unknown ligand type.")
        elif isinstance(ligand, Chem.Mol):
            self.load_ligand_from_mol(ligand)
        else:
            raise ValueError("Unknown ligand type.")

    def calculate_clashes(self) -> int:
        """Calculate the number of steric clashes between protein and ligand."""
        return [
            count_clashes(self.protein, mol, tollerance=self.clash_tolerance)
            for mol in self.ligands
        ]

    def calculate_strain_energy(self) -> float:
        """Calculate the strain energy of the ligand."""
        return [get_strain_energy(mol) for mol in self.ligands]

    def calculate_interactions(self) -> pd.DataFrame:
        """Calculate the interactions between the protein and the ligand."""
        return generate_interaction_df(self.protein, self.ligands)

    def calculate_rmsd(self, mol1: Chem.Mol, mol2: Chem.Mol) -> float:
        """Calculate the RMSD between two molecules

        For example, the RMSD between the generated ligand pose and one minimized by docking software.
        """
        return rmsd(mol1, mol2)

    def run(self) -> dict:
        """Execute all of PoseCheck suite and return the results.

        This method calculates the clashes, strain energy, and interactions
        between the protein and the ligand. The results are returned as a dictionary.

        Returns:
            dict: A dictionary containing the number of clashes, the strain energy,
                  and the interactions between the protein and the ligand.
        """
        clashes = self.calculate_clashes()
        strain = self.calculate_strain_energy()
        interactions = self.calculate_interactions()

        results = {"clashes": clashes, "strain": strain, "interactions": interactions}
        return results


if __name__ == "__main__":
    from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH

    """Test calculate_clashes method."""
    # Set the protein and ligand in the PoseCheck object
    pc = PoseCheck()
    pc.load_protein_from_pdb(EXAMPLE_PDB_PATH)
    pc.load_ligands_from_sdf(EXAMPLE_LIGAND_PATH)

    # Calculate the clashes
    clashes = pc.calculate_clashes()[0]

    # Check if the clashes are calculated correctly
    assert clashes == 2, "Clashes calculation is incorrect."

    """Test calculate_strain_energy method."""
    # Set the ligand in the PoseCheck object
    pc = PoseCheck()
    pc.load_ligands_from_sdf(EXAMPLE_LIGAND_PATH)

    # Calculate the strain energy
    strain_energy = pc.calculate_strain_energy()[0]

    # Check if the strain energy is calculated correctly
    assert round(strain_energy, 2) == 19.11

    """Test calculate_interactions method."""
    # Set the protein and ligand in the PoseCheck object
    pc = PoseCheck()
    pc.load_protein_from_pdb(EXAMPLE_PDB_PATH)
    pc.load_ligands_from_sdf(EXAMPLE_LIGAND_PATH)

    # Calculate the interactions
    interactions = pc.calculate_interactions()

    # Check if the interactions are calculated correctly
    assert interactions.shape[1] == 9, "Interactions calculation is incorrect."

    # Test calculate_rmsd method
    example_ligand = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)[0]

    rmsd_out = pc.calculate_rmsd(example_ligand, example_ligand)

    assert rmsd_out == 0.0, "RMSD calculation is incorrect."

posecheck/posecheck.py
cch1999/posecheck
import numpy as np
import rdkit.Chem as Chem
from rdkit import Chem
import datamol as dm


def rmsd(mol1: Chem.Mol, mol2: Chem.Mol) -> float:
    """
    Calculate the RMSD between two molecules.

    Args:
        mol1: The first molecule.
        mol2: The second molecule.

    Returns:
        The RMSD value between the two molecules.
    """

    if mol1 is None or mol2 is None:
        return np.nan

    mol1_coords = mol1.GetConformer().GetPositions()
    mol2_coords = mol2.GetConformer().GetPositions()

    try:
        rmsd = np.sqrt(np.mean((mol1_coords - mol2_coords) ** 2))
    except Exception as e:
        print(f"Error calculating RMSD: {e}")
        rmsd = np.nan

    return rmsd


def remove_radicals(mol: Chem.Mol, sanitize: bool = True) -> Chem.Mol:
    """Remove free radicals from a molecule."""

    # Iterate over all atoms in the molecule
    for atom in mol.GetAtoms():
        # Check if the current atom has any radical electrons
        num_radicals = atom.GetNumRadicalElectrons()
        if num_radicals > 0:
            # Saturate the atom with hydrogen atoms
            atom.SetNumExplicitHs(atom.GetNumExplicitHs() + num_radicals)
            atom.SetNumRadicalElectrons(0)

    if sanitize:
        dm.sanitize_mol(mol)

    return mol


def has_radicals(mol: Chem.Mol) -> bool:
    """Check if a molecule has any free radicals."""

    # Iterate over all atoms in the molecule
    for atom in mol.GetAtoms():
        # Check if the current atom has any radical electrons
        num_radicals = atom.GetNumRadicalElectrons()
        if num_radicals > 0:
            return True

    return False

posecheck/utils/chem.py
cch1999/posecheck
from typing import List, Union

import numpy as np
import rdkit.Chem as Chem
from rdkit import Chem

from posecheck.utils.chem import has_radicals


def count_clashes(prot: Chem.Mol, lig: Chem.Mol, tollerance: float = 0.5) -> int:
    """
    Counts the number of clashes between atoms in a protein and a ligand.

    Args:
        prot: RDKit Mol object representing the protein.
        lig: RDKit Mol object representing the ligand.
        tolerance: Distance tolerance for clash detection (default: 0.5).

    Returns:
        clashes: Number of clashes between the protein and the ligand.
    """

    # Check if the molecule has radicals
    assert not has_radicals(
        lig
    ), "Molecule has radicals, consider removing them first. (`posecheck.utils.chem.remove_radicals()`)"

    clashes = 0

    try:
        # Get the positions of atoms in the protein and ligand
        prot_pos = prot.GetConformer().GetPositions()
        lig_pos = lig.GetConformer().GetPositions()

        pt = Chem.GetPeriodicTable()

        # Get the number of atoms in the protein and ligand
        num_prot_atoms = prot.GetNumAtoms()
        num_lig_atoms = lig.GetNumAtoms()

        # Calculate the Euclidean distances between all atom pairs in the protein and ligand
        dists = np.linalg.norm(
            prot_pos[:, np.newaxis, :] - lig_pos[np.newaxis, :, :], axis=-1
        )

        # Iterate over the ligand atoms
        for i in range(num_lig_atoms):
            lig_vdw = pt.GetRvdw(lig.GetAtomWithIdx(i).GetAtomicNum())

            # Iterate over the protein atoms
            for j in range(num_prot_atoms):
                prot_vdw = pt.GetRvdw(prot.GetAtomWithIdx(j).GetAtomicNum())

                # Check for clash by comparing the distances with tolerance
                if dists[j, i] + tollerance < lig_vdw + prot_vdw:
                    clashes += 1

    except AttributeError:
        raise ValueError(
            "Invalid input molecules. Please provide valid RDKit Mol objects."
        )

    return clashes


#### Helper functions


def count_clashes_list(
    prot: Chem.Mol,
    ligs: List[Union[Chem.Mol, str]],
    target: Union[str, None] = None,
    tollerance: float = 0.5,
) -> List[dict]:
    """
    Counts the number of clashes between atoms in a protein and a list of ligands.

    Args:
        prot: RDKit Mol object representing the protein.
        ligs: List of RDKit Mol objects or SMILES strings representing the ligands.
        target: Target identifier associated with the ligands (default: None).
        tolerance: Distance tolerance for clash detection (default: 0.5).

    Returns:
        clashes: List of dictionaries containing clash information for each ligand.
                 Each dictionary contains the following keys: 'mol' (RDKit Mol object),
                 'clashes' (number of clashes), and 'target' (target identifier).
    """
    clashes = []

    # Iterate over the ligands
    for lig in ligs:
        try:
            # Create RDKit Mol object from SMILES string or existing Mol object
            if isinstance(lig, str):
                lig = Chem.MolFromSmiles(lig)
            else:
                lig = Chem.Mol(lig)

            # Count clashes between protein and ligand
            lig_clashes = count_clashes(prot, lig, tollerance)

            # Append clash information to the clashes list
            clashes.append({"mol": lig, "clashes": lig_clashes, "target": target})
        except Exception as e:
            # Handle errors by appending a dictionary with NaN clashes
            clashes.append({"mol": lig, "clashes": np.nan, "target": target})

    return clashes


if __name__ == "__main__":
    from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH
    from posecheck.utils.loading import load_mols_from_sdf, load_protein_from_pdb

    # Example molecules
    prot = load_protein_from_pdb(EXAMPLE_PDB_PATH)
    lig = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)[0]

    # Count clashes between protein and ligand
    clashes = count_clashes(prot, lig)

    # Print the number of clashes
    print(f"Number of clashes: {clashes}")

    assert clashes == 2

posecheck/utils/clashes.py
cch1999/posecheck
import os

# get directory of this file
DIR = os.path.dirname(os.path.abspath(__file__))
# get directory of project and data
DATA_DIR = os.path.join(DIR, "../../data/")
PROJECT_DIR = os.path.join(DIR, "../../")

CROSSOCKED_PDB_DIR = os.path.join(DATA_DIR, "crossdocked/")
CROSSOCKED_PDBQT_DIR = os.path.join(DATA_DIR, "crossdocked_pdbqt/")

PROTEINS_PATH = os.path.join(DATA_DIR, "proteins.pt")

# Example paths

EXAMPLE_PDB_PATH = os.path.join(DATA_DIR, "examples/1a2g.pdb")
EXAMPLE_LIGAND_PATH = os.path.join(DATA_DIR, "examples/1a2g_ligand.sdf")

# Constants collected from the codebase
FORCEFIELD = "uff"
ADD_COORDS = True
ROUND_DIGITS = 2

# PATHS
# REDUCE_PATH = os.path.join(PROJECT_DIR, "bin/reduce")
REDUCE_PATH = "reduce"
SPLIT_PATH = os.path.join(PROJECT_DIR, "data/crossdocked_split_by_name.pkl")

# Docking params
DOCKING_METHOD = "smina"  # TODO add support for vina and qvina2
# SMINA_PATH = '/home/cch57/projects/poses_benchmark/smina/smina.static'
SMINA_PATH = os.path.join(PROJECT_DIR, "bin/smina.static")
EXHAUSTIVENESS = 8
BOX_SIZE = 25

posecheck/utils/constants.py
cch1999/posecheck
import os
from typing import Union

import datamol as dm
import numpy as np
from rdkit import Chem

from posecheck.utils.constants import BOX_SIZE, SMINA_PATH
from posecheck.utils.docking_utils import *

# TODO add this to .env file
SMINA_PATH = "/home/cch57/projects/poses_benchmark/smina/smina.static"


class SMINA(object):
    def __init__(
        self,
        executable=SMINA_PATH,
        cpu=None,
    ):
        """Wrapper class for smina docking software.

        Args:
            executable (str): Path to smina executable.
            cpu (int, optional): Number of CPU cores to use. If None, all available cores will be used (recommended).
        """
        self.executable = executable
        self.minimize = False
        self.score_only = False
        self.cpu = cpu

    def set_receptor(
        self,
        receptor_path: Union[str, os.PathLike],
        centre: Union[str, Chem.Mol, np.ndarray] = "pocket",
        size: int = BOX_SIZE,
    ) -> None:
        """
        Set the receptor for docking.

        Args:
            receptor_path (Union[str, os.PathLike]): Path to the receptor file.
            centre (Union[str, Chem.Mol, np.ndarray]): The centre of the docking box.
                If 'pocket', the centre is calculated from the PDBQT file.
                If Chem.Mol, the centre is calculated from the RDKit molecule.
                If np.ndarray, the centre is already specified as a numpy array.
                If str, the centre is read from an SDF file.
            size (int): The size of the docking box. Default is 25.

        Raises:
            ValueError: If the centre type is not one of the specified types.
        """

        self.receptor = receptor_path

        if centre == "pocket":
            # Get the receptor centre from the PDBQT file with rdkit
            receptor_mol = read_pdbqt_receptor(receptor_path)
            centre = get_centroid_ob(receptor_mol)
        elif type(centre) == Chem.Mol:
            # Get the receptor centre from the rdkit mol
            centre = get_centroid_rdmol(centre)
        elif type(centre) == np.ndarray:
            # Centre is already a numpy array
            pass
        elif type(centre) == str:  # TODO add support for PDBQT
            # get from SDF
            try:
                mol = dm.read_sdf(centre)[0]
                centre = get_centroid_rdmol(mol)
            except:
                raise ValueError(
                    f"Could not read centre from file {centre}. Make sure it is a valid SDF file."
                )
        else:
            raise ValueError(
                f'Invalid centre type {type(centre)}. Must be one of: "pocket", Chem.Mol, np.ndarray, or path to sdf.'
            )

        # Set the centre and size of the docking box
        self.centre_x = centre[0]
        self.centre_y = centre[1]
        self.centre_z = centre[2]

        self.size_x = size
        self.size_y = size
        self.size_z = size

    def set_ligand(self, ligand_path):
        raise NotImplementedError(
            "Use set_ligand_from_mol instead, centering is not implemented for this method."
        )
        # self.ligand = ligand_path

    def set_ligand_from_mol(self, mol: Chem.Mol):
        self.task_id = np.random.randint(0, 1000000)
        self.ligand = f"tmp_{self.task_id}.pdbqt"

        rdkit_mol_to_pdbqt(mol, self.ligand)

    def clear_ligand(self):
        """Remove the ligand file. Important to do this after each docking run."""
        if os.path.exists(self.ligand):
            os.remove(self.ligand)

    def score_pose(self):
        self.score_only = True
        mol, out = self.run()
        self.score_only = False

        out = parse_smina_output_score(out)

        if len(mol) == 1:
            out["mol"] = mol[0]
        else:
            out["mol"] = None

        return out

    def minimize_pose(self):
        self.minimize = True
        mol, out = self.run()
        self.minimize = False

        out = parse_smina_output_minimize(out)

        if len(mol) == 1:
            out["mol"] = mol[0]
        else:
            out["mol"] = None

        return out

    def redock(self):
        mols, out = self.run()
        out = parse_smina_output_docking(out)

        out["mols"] = mols

        return out

    def calculate_all(self, mol=None):
        if mol is not None:
            self.set_ligand_from_mol(mol)

        try:
            score_only = self.score_pose()
            minimized = self.minimize_pose()
            redocked = self.redock()

            self.clear_ligand()

            return {
                "score_only": score_only,
                "minimized": minimized,
                "redocked": redocked,
            }

        except:
            self.clear_ligand()
            return {"score_only": None, "minimized": None, "redocked": None}

    def run(self):
        """Main function to run smina.
        Should be called from one of the other functions.
        """

        tmp_out = f"tmp_out_{self.task_id}.pdbqt"

        # Command over multiple lines for readability
        command = (
            f"{self.executable} -r {self.receptor} -l {self.ligand} "
            f"--center_x {self.centre_x} --center_y {self.centre_y} --center_z {self.centre_z} "
            f"--size_x {self.size_x} --size_y {self.size_y} --size_z {self.size_z} "
            f"-o {tmp_out} "
        )

        if not self.minimize and not self.score_only:
            command += f"--exhaustiveness 8"

        if self.minimize:
            command += "--minimize "
        if self.score_only:
            command += "--score_only "

        if self.cpu:
            command += f"--cpu {self.cpu} "  # only do if you want use less than the max number of cores on your machine

        # Perform docking/minimization
        # print(command)
        # os.system(command)
        out = os.popen(command).read()

        out_mols = pdbqt_to_rdkit_mols(tmp_out)
        os.remove(tmp_out)

        return out_mols, out

posecheck/utils/docking.py
cch1999/posecheck
import os
from typing import Dict

import numpy as np
import openbabel
import openbabel as ob
import rdkit.Chem as Chem
from rdkit import Chem

#### AutoDock Vina Parsing Functions ####


def parse_smina_output_minimize(out) -> Dict[str, float]:
    """Parse the output of smina --minimize.

    Parameters
    ----------
    out : str
        The output of smina --minimize.

    Returns
    -------
    dict of str: float
    """

    lines = out.split("\n")

    affinity = float(lines[-5].split(" ")[1])
    rmsd = float(lines[-4].split(" ")[1])

    return {"min_affinity": affinity, "min_rmsd": rmsd}


def parse_smina_output_score(out: str) -> dict:
    """
    Parses the output of smina and returns the affinity score, intramolecular energy and energy terms.

    Arguments:
        out: The output of smina as a string.

    Returns:
        A dictionary containing the affinity score, intramolecular energy and energy terms.
    """
    # Split the output into lines
    lines = out.split("\n")

    # Extract the affinity score
    affinity = float(lines[-7].split(" ")[1])

    # Extract the intramolecular energy
    intramolecular_energy = float(lines[-6].split(" ")[-1])

    # Extract the energy terms
    energy_terms = [float(lines[-4].split(" ")[i + 2]) for i in range(6)]

    return {
        "affinity": affinity,
        "intramolecular_energy": intramolecular_energy,
        "energy_terms": energy_terms,
    }


def parse_smina_output_docking(out: str) -> dict:
    """
    Parses Smina output

    Parameters
    ----------
    out : str
        Smina output string

    Returns
    -------
    dict
        Dictionary containing the affinities
    """
    # Check if there is an output
    if "-----+------------+----------+----------" not in out:
        return None

    # Split output into lines
    out_split = out.splitlines()
    # Find the line with the best affinity
    best_idx = out_split.index("-----+------------+----------+----------") + 1

    # Get all of the affinities
    affinities = []

    # Loop poses until the 'Refine' line is reached
    while True:
        line = out_split[best_idx]
        if line.split()[0] == "Refine":
            break
        else:
            affinities.append(float(line.split()[1]))
            best_idx += 1

    return {"affinities": affinities}


#### Open Babel Parsing Functions For Receptor Centroid ####


def read_pdbqt_receptor(file_path):
    """Read a PDBQT file and return an Open Babel molecule object.
    This is done to get the centroid of the receptor.
    """
    # Create an Open Babel molecule object
    ob_mol = openbabel.OBMol()

    # Create an Open Babel file format object for PDBQT
    file_format = openbabel.OBConversion()
    file_format.SetInAndOutFormats("pdbqt", "pdbqt")

    # Read the PDBQT file
    if file_format.ReadFile(ob_mol, file_path):
        return ob_mol
    else:
        raise ValueError(f"Error reading PDBQT file: {file_path}")


def get_coordinates_ob(molecule):
    coordinates = []
    for atom in openbabel.OBMolAtomIter(molecule):
        x, y, z = atom.GetX(), atom.GetY(), atom.GetZ()
        coordinates.append((x, y, z))
    return coordinates


def get_centroid_ob(molecule):
    """Get centroid of molecule"""
    coordinates = get_coordinates_ob(molecule)
    centroid = np.mean(coordinates, axis=0)
    return centroid


def get_centroid_rdmol(mol: Chem.Mol) -> np.ndarray:
    """Calculate the centroid of a molecule.

    Args:
        mol (Chem.Mol): RDKit molecule object.

    Returns:
        np.ndarray: The centroid of the molecule.
    """
    conf = mol.GetConformer()
    positions = conf.GetPositions()
    centroid = np.mean(positions, axis=0)
    return centroid


def update_coordinates(mol, new_coords):
    """
    Update the coordinates of an RDKit molecule.

    Parameters:
    - mol: RDKit molecule object
    - new_coords: 2D list or array of shape (num_atoms, 3) containing the new coordinates

    Returns:
    - RDKit molecule with updated coordinates
    """

    if not mol.GetNumConformers():
        raise ValueError("The molecule has no conformers.")

    conf = mol.GetConformer()

    for idx, (x, y, z) in enumerate(new_coords):
        conf.SetAtomPosition(idx, (x, y, z))

    return mol


#### Read and Write PDBQT Files to RDKIT Mols ####


def rdkit_mol_to_pdbqt(rdkit_mol, pdbqt_path):
    """Save an RDKit molecule to a PDBQT file using an intermediate PDB file."""

    # Create a temporary PDB file path with a random name
    temp_pdb_path = f"temp_molecule_{np.random.randint(1000000)}.pdb"

    # print(rdkit_mol.GetConformer().GetPositions()[:4])
    # TODO REMOVE

    # AllChem.EmbedMolecule(rdkit_mol, AllChem.ETKDG())
    # ! NOTE careful! this added random coords

    # Write RDKit molecule to the temporary PDB file
    writer = Chem.PDBWriter(temp_pdb_path)
    writer.write(rdkit_mol)
    writer.close()

    # Set up Open Babel objects
    mol = ob.OBMol()
    obConversion = ob.OBConversion()
    obConversion.SetInAndOutFormats("pdb", "pdbqt")

    # Read the temporary PDB file with Open Babel
    if not obConversion.ReadFile(mol, temp_pdb_path):
        raise ValueError(f"Failed to read the temporary PDB file: {temp_pdb_path}")

    # Write the molecule to a PDBQT file
    if not obConversion.WriteFile(mol, pdbqt_path):
        raise ValueError(f"Failed to write to the PDBQT file: {pdbqt_path}")

    # Clean up temporary PDB file
    os.remove(temp_pdb_path)


def pdbqt_to_rdkit_mols(pdbqt_path):
    """Convert all molecules in a PDBQT file to a list of RDKit molecules."""

    # Set up Open Babel objects
    mol = ob.OBMol()
    obConversion = ob.OBConversion()
    obConversion.SetInAndOutFormats("pdbqt", "pdb")

    # Ensure the file can be opened
    if not obConversion.OpenInAndOutFiles(pdbqt_path, ""):
        raise ValueError(f"Failed to read the PDBQT file: {pdbqt_path}")

    rdkit_mols = []
    # Loop over all molecules in the file
    not_end_of_file = True
    while not_end_of_file:
        not_end_of_file = obConversion.Read(mol)
        # Convert to PDB format
        pdb_block = obConversion.WriteString(mol)

        # Convert PDB block to RDKit molecule
        rdkit_mol = Chem.MolFromPDBBlock(pdb_block)

        if rdkit_mol is not None:
            rdkit_mols.append(rdkit_mol)

        # Clear the molecule object for the next iteration
        mol.Clear()

    if not rdkit_mols:
        raise ValueError("No molecules could be successfully converted.")

    # Remove the last molecule, which is always empty
    return rdkit_mols[:-1]

posecheck/utils/docking_utils.py
cch1999/posecheck
from typing import List, Optional

import pandas as pd
import prolif as plf
from rdkit import DataStructs


def generate_interaction_df(prot: plf.Molecule, lig: plf.Molecule) -> pd.DataFrame:
    """
    Generate a DataFrame with all interactions between protein and ligand.

    Args:
        prot: A protein molecule of type plf.Molecule.
        lig: A ligand molecule of type plf.Molecule.

    Returns:
        A DataFrame representing all interactions between the protein and ligand.
    """
    fp = plf.Fingerprint()
    fp.run_from_iterable(lig, prot)
    df = fp.to_dataframe()
    return df


def merge_interaction_dfs(ref_df: pd.DataFrame, df: pd.DataFrame) -> pd.DataFrame:
    """
    Merge two interaction DataFrames into one.

    Args:
        ref_df: The reference interaction DataFrame of type pd.DataFrame.
        df: The interaction DataFrame to merge with the reference of type pd.DataFrame.

    Returns:
        A merged DataFrame containing interactions from both input DataFrames.
    """

    ref_df.rename({0: "ref"}, inplace=True)
    # drop the ligand level on both dataframes
    ref_df.columns = ref_df.columns.droplevel(0)
    df.columns = df.columns.droplevel(0)
    # concatenate and sort columns
    df = (
        pd.concat([ref_df, df])
        .fillna(False)
        .sort_index(
            axis=1,
            level=0,
            key=lambda index: [plf.ResidueId.from_string(x) for x in index],
        )
    )
    return df


def calculate_interaction_similarity(df: Optional[pd.DataFrame]) -> List[float]:
    """
    Calculate the Tanimoto similarity between the reference ligand and all other ligands.

    Args:
        df: A DataFrame containing interaction data. If None, an empty list is returned.

    Returns:
        A list of Tanimoto similarity values between the reference ligand and all other ligands.
    """

    if df is None:
        return []

    bvs = plf.to_bitvectors(df)
    similarities = []
    for i, bv in enumerate(bvs[1:]):
        tc = DataStructs.TanimotoSimilarity(bvs[0], bv)
        # print(f"{i}: {tc:.3f}")
        similarities.append(round(tc, 2))
    return similarities


if __name__ == "__main__":
    from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH
    from posecheck.utils.loading import load_mols_from_sdf, load_protein_from_pdb

    # Example usage of the functions
    prot = load_protein_from_pdb(EXAMPLE_PDB_PATH)
    lig = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)

    interaction_df = generate_interaction_df(prot, lig)
    print(interaction_df)

    # Note that the reference ligand is the first ligand in the DataFrame
    # This is usually done to compare generated ligands to a reference ligand
    # In this case, we are comparing the reference ligand to itself as an example
    merged_df = merge_interaction_dfs(interaction_df, interaction_df)
    print(merged_df)

    similarities = calculate_interaction_similarity(merged_df)
    print(similarities)

posecheck/utils/interactions.py
cch1999/posecheck
import os
import pickle
import subprocess
from functools import lru_cache
from typing import List, Union

import datamol as dm
import MDAnalysis as mda
import openbabel as ob
import prolif as plf
import rdkit.Chem as Chem
from rdkit import Chem

from posecheck.utils.chem import remove_radicals
from posecheck.utils.constants import REDUCE_PATH, SPLIT_PATH


def load_splits_crossdocked(
    split_path: str = SPLIT_PATH,
):
    """Load the train, val, and test splits."""
    with open(split_path, "rb") as f:
        data = pickle.load(f)
    return data


def get_ids_to_pockets(
    split_path: str = SPLIT_PATH,
):
    """Get the names of the pockets in the CrossDocked test dataset."""
    test_set = load_splits_crossdocked(split_path)["test"]

    pdb_names = [
        pdb_path[0].split("/")[1].split(".")[0].replace("_", "-")
        for pdb_path in test_set
    ]

    return pdb_names


def load_protein_prolif(protein_path: str):
    """Load protein from PDB file using MDAnalysis
    and convert to plf.Molecule. Assumes hydrogens are present."""
    prot = mda.Universe(protein_path)
    prot = plf.Molecule.from_mda(prot, NoImplicit=False)
    return prot


@lru_cache(maxsize=100)
def load_protein_from_pdb(pdb_path: str, reduce_path: str = REDUCE_PATH):
    """Load protein from PDB file, add hydrogens, and convert it to a prolif.Molecule.

    Args:
        pdb_path (str): The path to the PDB file.
        reduce_path (str, optional): The path to the reduce executable. Defaults to REDUCE_PATH.

    Returns:
        plf.Molecule: The loaded protein as a prolif.Molecule.
    """
    tmp_path = pdb_path.split(".pdb")[0] + "_tmp.pdb"

    # Call reduce to make tmp PDB with waters
    reduce_command = f"{reduce_path} -NOFLIP  {pdb_path} -Quiet > {tmp_path}"
    subprocess.run(reduce_command, shell=True)

    # Load the protein from the temporary PDB file
    prot = load_protein_prolif(tmp_path)
    os.remove(tmp_path)

    return prot


def get_pdbqt_mol(pdbqt_block: str) -> Chem.Mol:
    """Convert pdbqt block to rdkit mol by converting with openbabel"""
    # write pdbqt file
    with open("test_pdbqt.pdbqt", "w") as f:
        f.write(pdbqt_block)

    # read pdbqt file from autodock
    mol = ob.OBMol()
    obConversion = ob.OBConversion()
    obConversion.SetInAndOutFormats("pdbqt", "pdb")
    obConversion.ReadFile(mol, "test_pdbqt.pdbqt")

    # convert to RDKIT
    mol = Chem.MolFromPDBBlock(obConversion.WriteString(mol))

    # remove tmp file
    os.remove("test_pdbqt.pdbqt")

    return mol


def read_pdbqt(pdbqt_file: str) -> Chem.Mol:
    """Read pdbqt file from autodock and convert to rdkit mol.

    Args:
        pdbqt_file (str): path to pdbqt file

    Returns:
        Chem.Mol: rdkit mol
    """
    with open(pdbqt_file, "r") as f:
        pdbqt_block = f.read()
    return get_pdbqt_mol(pdbqt_block)


def load_sdf_prolif(sdf_path: str) -> plf.Molecule:
    """Load ligand from an SDF file and convert it to a prolif.Molecule.

    Args:
        sdf_path (str): Path to the SDF file.

    Returns:
        plf.Molecule: The loaded ligand as a prolif.Molecule.
    """
    return plf.sdf_supplier(sdf_path)


def load_mols_from_sdf(sdf_path: str, add_hs: bool = True) -> Union[plf.Molecule, List]:
    """Load ligand from an SDF file, add hydrogens, and convert it to a prolif.Molecule.

    Args:
        sdf_path (str): Path to the SDF file.
        add_hs (bool, optional): Whether to add hydrogens. Defaults to True.

    Returns:
        Union[plf.Molecule, List]: The loaded ligand as a prolif.Molecule, or an empty list if no molecule could be loaded.
    """
    tmp_path = sdf_path.split(".sdf")[0] + "_tmp.sdf"

    # Load molecules from the SDF file
    mols = dm.read_sdf(sdf_path)

    # Remove radicals from the molecules
    mols = [remove_radicals(mol) for mol in mols]

    # Add hydrogens to the molecules
    if add_hs:
        mols = [Chem.AddHs(m, addCoords=True) for m in mols if m is not None]

    if len(mols) == 0:
        return []

    # Write the molecules to a temporary file
    dm.to_sdf(mols, tmp_path)
    ligs = load_sdf_prolif(tmp_path)
    os.remove(tmp_path)

    # Turn into list
    ligs = list(ligs)

    return ligs


def load_mols_from_rdkit(
    mol_list: List[Chem.Mol], add_hs: bool = True
) -> Union[plf.Molecule, List]:
    """Load ligand from an RDKit molecule, add hydrogens, and convert it to a prolif.Molecule.

    Processes in the same way as src.utils.loading.load_mols_from_sdf but takes an RDKit molecule as input.

    Args:
        rdkit_mol (Chem.Mol): RDKit molecule.
        add_hs (bool, optional): Whether to add hydrogens. Defaults to True.

    Returns:
        Union[plf.Molecule, List]: The loaded ligand as a prolif.Molecule, or an empty list if no molecule could be loaded.
    """

    # Convert single molecule to list
    if isinstance(mol_list, Chem.Mol):
        print(
            "Converting single molecule to list, to be consistent with load_mols_from_sdf"
        )
        mol_list = [mol_list]

    # Add hydrogens to the molecules
    if add_hs:
        mol_list = [Chem.AddHs(m, addCoords=True) for m in mol_list if m is not None]

    # Check if any molecules were loaded
    if len(mol_list) == 0:
        print("No molecules loaded, returning empty list")
        return []

    return [plf.Molecule.from_rdkit(mol) for mol in mol_list]


if __name__ == "__main__":
    from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH

    prot = load_protein_from_pdb(EXAMPLE_PDB_PATH)
    lig = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)[0]

    mol = dm.read_sdf(EXAMPLE_LIGAND_PATH)[0]
    lig = load_mols_from_rdkit(mol)

    # Crossdocked Name
    names = get_ids_to_pockets()
    print(f"Number of crossdocked pockets in test set: {len(names)}")

posecheck/utils/loading.py
cch1999/posecheck
from copy import deepcopy

import numpy as np
import rdkit.Chem as Chem
from datamol.conformers._conformers import _get_ff
from rdkit import Chem
from rdkit.Chem import AllChem

from posecheck.utils.chem import has_radicals


def calculate_energy(
    mol: Chem.Mol, forcefield: str = "UFF", add_hs: bool = True
) -> float:
    """
    Calculates the energy of a molecule using a force field.

    Args:
        mol: RDKit Mol object representing the molecule.
        forcefield: Force field to use for energy calculation (default: "UFF").
        add_hs: Whether to add hydrogens to the molecule (default: True).

    Returns:
        energy: Calculated energy of the molecule (rounded to 2 decimal places).
                Returns NaN if energy calculation fails.
    """
    mol = Chem.Mol(mol)  # Make a deep copy of the molecule

    if add_hs:
        mol = Chem.AddHs(mol, addCoords=True)

    try:
        ff = _get_ff(mol, forcefield=forcefield)
    except Exception:
        return np.nan

    energy = ff.CalcEnergy()

    return round(energy, 2)


def relax_mol(mol: Chem.Mol) -> Chem.Mol:
    """Relax a molecule by adding hydrogens, embedding it, and optimizing it.

    Args:
        mol (Chem.Mol): The molecule to relax.

    Returns:
        Chem.Mol: The relaxed molecule.
    """

    # if the molecule is None, return None
    if mol is None:
        return None

    # Incase ring info is not present
    Chem.GetSSSR(mol)  # SSSR: Smallest Set of Smallest Rings

    # make a copy of the molecule
    mol = deepcopy(mol)

    # add hydrogens
    mol = Chem.AddHs(mol, addCoords=True)

    # embed the molecule
    AllChem.EmbedMolecule(mol, randomSeed=0xF00D)

    # optimize the molecule
    AllChem.UFFOptimizeMolecule(mol)

    # return the molecule
    return mol


def get_strain_energy(mol: Chem.Mol) -> float:
    """
    Calculates the strain energy of a molecule.

    Strain energy is defined as the difference between the energy of a molecule
    and the energy of the same molecule in a relaxed geometry.

    Parameters
    ----------
    mol : rdkit.Chem.rdchem.Mol
        Molecule.

    Returns
    -------
    float
        Strain energy.

    """

    # Check if the molecule has radicals
    assert not has_radicals(
        mol
    ), "Molecule has radicals, consider removing them first. (`posecheck.utils.chem.remove_radicals()`)"

    try:
        return calculate_energy(mol) - calculate_energy(relax_mol(mol))
    except:
        return np.nan


if __name__ == "__main__":
    from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH
    from posecheck.utils.loading import load_mols_from_sdf, load_protein_from_pdb

    # Example molecules
    prot = load_protein_from_pdb(EXAMPLE_PDB_PATH)
    lig = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)[0]

    # Calculate strain
    strain = get_strain_energy(lig)

    print(f"Strain energy: {strain}")

    assert round(strain, 2) == 19.11

posecheck/utils/strain.py
cch1999/posecheck
from setuptools import setup, find_packages

setup(
    name="posecheck",
    version="0.1",
    description="A library for pose estimation benchmarks",
    author="Charles Harris",
    author_email="cch57@cam.ac.uk",
    packages=find_packages(include=["posecheck", "posecheck.*"]),
    # install_requires=[
    #    "numpy>=1.19.5",
    #    "tensorflow>=2.6",
    #    # Add other dependencies here
    # ],
    classifiers=[
        "Development Status :: 3 - Alpha",
        "Intended Audience :: Science/Research",
        "Topic :: Scientific/Engineering :: Bio-Informatics",
        "License :: OSI Approved :: MIT License",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.6",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
    ],
)

setup.py
cch1999/posecheck
import unittest

import datamol as dm
import rdkit.Chem as Chem
import rdkit.Chem.AllChem as AllChem
import numpy as np

from posecheck.utils.chem import rmsd, has_radicals, remove_radicals
from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH
from posecheck.utils.strain import calculate_energy, get_strain_energy, relax_mol


class TestChem(unittest.TestCase):
    def test_rmsd(self):
        mol1 = dm.read_sdf(EXAMPLE_LIGAND_PATH)[0]
        mol2 = dm.read_sdf(EXAMPLE_LIGAND_PATH)[0]
        result = rmsd(mol1, mol2)
        self.assertEqual(result, 0.0)

    def test_remove_radicals(self):
        """Test the remove_radicals function."""
        # Create a molecule with free radicals
        mol = Chem.MolFromSmiles("[O]")

        # Check if the molecule has radicals before removal
        self.assertTrue(has_radicals(mol))

        # Remove the radicals
        mol = remove_radicals(mol)

        # Check if the molecule has radicals after removal
        self.assertFalse(has_radicals(mol))

    def test_remove_radicals2(self):
        # Define a test molecule

        # smile for large molecule
        smiles = "CC(CC1=CC=C(C=C1)C(=O)O)C1=CC=C(C=C1)C(=O)O"

        mol = Chem.MolFromSmiles(smiles)
        # add generate conformer using UFF
        AllChem.EmbedMolecule(mol, randomSeed=0xF00D)

        # set radical
        for i, atom in enumerate(mol.GetAtoms()):
            if atom.GetSymbol() == "O":
                if atom.GetSymbol() == "C":
                    atom.SetNumRadicalElectrons(1)
                    break

        mol = remove_radicals(mol)

        # Calculate the energy using the UFF force field
        energy = get_strain_energy(mol)

        # Assert that the calculated energy is not NaN
        self.assertFalse(np.isnan(energy))

    def test_has_radicals(self):
        """Test the has_radicals function."""
        # Create a molecule without radicals
        mol = Chem.MolFromSmiles("CC")

        # Check if the molecule has radicals
        self.assertFalse(has_radicals(mol))

        # Create a molecule with radicals
        mol = Chem.MolFromSmiles("[O]")

        # Check if the molecule has radicals
        self.assertTrue(has_radicals(mol))


if __name__ == "__main__":
    unittest.main()

tests/test_chem.py
cch1999/posecheck
import unittest

from posecheck.utils.clashes import count_clashes, count_clashes_list
from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH
from posecheck.utils.loading import load_mols_from_sdf, load_protein_from_pdb


class ClashesTest(unittest.TestCase):
    def setUp(self):
        # Define a test protein and ligand
        self.prot = load_protein_from_pdb(EXAMPLE_PDB_PATH)
        self.lig = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)[0]

    def test_count_clashes(self):
        # Count clashes between protein and ligand
        clashes = count_clashes(self.prot, self.lig)

        # Assert that the number of clashes is correct
        self.assertEqual(clashes, 2)

    def test_count_clashes_list(self):
        # Count clashes between protein and ligands
        clashes_list = count_clashes_list(self.prot, [self.lig, self.lig])

        # Assert that the number of clashes for each ligand is correct
        self.assertEqual(clashes_list[0]["clashes"], 2)
        self.assertEqual(clashes_list[1]["clashes"], 2)


if __name__ == "__main__":
    unittest.main()

tests/test_clashes.py
cch1999/posecheck
import unittest

import pandas as pd

from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH
from posecheck.utils.interactions import (calculate_interaction_similarity,
                                    generate_interaction_df,
                                    merge_interaction_dfs)
from posecheck.utils.loading import load_mols_from_sdf, load_protein_from_pdb


class TestInteractions(unittest.TestCase):
    def setUp(self):
        self.prot = load_protein_from_pdb(EXAMPLE_PDB_PATH)
        self.lig = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)

    def test_generate_interaction_df(self):
        interaction_df = generate_interaction_df(self.prot, self.lig)
        self.assertIsInstance(interaction_df, pd.DataFrame)
        # Add more assertions as needed

    def test_merge_interaction_dfs(self):
        interaction_df = generate_interaction_df(self.prot, self.lig)
        merged_df = merge_interaction_dfs(interaction_df, interaction_df)
        self.assertIsInstance(merged_df, pd.DataFrame)
        # Add more assertions as needed

    def test_calculate_interaction_similarity(self):
        interaction_df = generate_interaction_df(self.prot, self.lig)
        merged_df = merge_interaction_dfs(interaction_df, interaction_df)
        similarities = calculate_interaction_similarity(merged_df)
        self.assertEqual(similarities, [1.0])
        # Add more assertions as needed


if __name__ == "__main__":
    unittest.main()

tests/test_interactions.py
cch1999/posecheck
import unittest

import datamol as dm

from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH
from posecheck.utils.loading import (get_ids_to_pockets, load_mols_from_rdkit,
                               load_mols_from_sdf, load_protein_from_pdb)


class LoadingTest(unittest.TestCase):
    def test_load_protein_from_pdb(self):
        prot = load_protein_from_pdb(EXAMPLE_PDB_PATH)
        self.assertIsNotNone(prot)

    def test_load_mols_from_sdf(self):
        lig = load_mols_from_sdf(EXAMPLE_LIGAND_PATH)[0]
        self.assertIsNotNone(lig)

    def test_load_mols_from_rdkit(self):
        mol = dm.read_sdf(EXAMPLE_LIGAND_PATH)[0]
        lig = load_mols_from_rdkit(mol)
        self.assertIsNotNone(lig)
        self.assertIsInstance(lig, list)

    def test_get_ids_to_pockets(self):
        names = get_ids_to_pockets()
        self.assertIsNotNone(names)


if __name__ == "__main__":
    unittest.main()

tests/test_loading.py
cch1999/posecheck
import unittest

import datamol as dm

from posecheck.posecheck import PoseCheck
from posecheck.utils.constants import EXAMPLE_LIGAND_PATH, EXAMPLE_PDB_PATH


class TestPoseCheck(unittest.TestCase):
    def setUp(self):
        self.pc = PoseCheck()
        self.pc.load_protein_from_pdb(EXAMPLE_PDB_PATH)
        self.pc.load_ligands_from_sdf(EXAMPLE_LIGAND_PATH)

    def test_load_ligands_from_rdmol(self):
        mol = dm.read_sdf(EXAMPLE_LIGAND_PATH)[0]
        pc_test = PoseCheck()
        pc_test.load_ligands_from_mols(mol)

    def test_calculate_clashes(self):
        clashes = self.pc.calculate_clashes()[0]
        self.assertEqual(clashes, 2, "Clashes calculation is incorrect.")

    def test_calculate_strain_energy(self):
        strain_energy = self.pc.calculate_strain_energy()[0]
        self.assertAlmostEqual(strain_energy, 19.11, places=2)

    def test_calculate_interactions(self):
        interactions = self.pc.calculate_interactions()
        self.assertEqual(
            interactions.shape[1], 9, "Interactions calculation is incorrect."
        )

    def test_calculate_rmsd(self):
        example_ligand = dm.read_sdf(EXAMPLE_LIGAND_PATH)[0]
        rmsd_out = self.pc.calculate_rmsd(example_ligand, example_ligand)
        self.assertEqual(rmsd_out, 0.0, "RMSD calculation is incorrect.")

tests/test_posecheck.py
cch1999/posecheck
import os
import subprocess
import unittest

from posecheck.utils.constants import EXAMPLE_PDB_PATH, REDUCE_PATH


class TestReduce(unittest.TestCase):
    def test_reduce_exists(self):
        command = REDUCE_PATH
        exit_response = subprocess.run(
            command, capture_output=True, text=True, shell=True
        )
        self.assertEqual(exit_response.returncode, 1)

    def test_reduce_executable(self):
        command = REDUCE_PATH + f" -NOFLIP {EXAMPLE_PDB_PATH}"
        exit_response = subprocess.run(
            command, capture_output=True, text=True, shell=True
        )
        self.assertEqual(exit_response.returncode, 0)

tests/test_reduce.py
cch1999/posecheck
import unittest

import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem

from posecheck.utils.strain import calculate_energy, get_strain_energy, relax_mol


class EnergyCalculationTest(unittest.TestCase):
    def test_calculate_energy(self):
        # Define a test molecule

        # smile for large molecule
        smiles = "CC(CC1=CC=C(C=C1)C(=O)O)C1=CC=C(C=C1)C(=O)O"

        mol = Chem.MolFromSmiles(smiles)
        # add generate conformer using UFF
        AllChem.EmbedMolecule(mol, randomSeed=0xF00D)

        # Calculate the energy using the UFF force field
        energy = calculate_energy(mol, forcefield="UFF", add_hs=True)

        # Assert that the calculated energy is not NaN
        self.assertFalse(np.isnan(energy))


class MoleculeRelaxationTest(unittest.TestCase):
    def test_relax_mol(self):
        # Define a test molecule
        smiles = "CCO"
        mol = Chem.MolFromSmiles(smiles)

        # Relax the molecule
        relaxed_mol = relax_mol(mol)

        # Assert that the relaxed molecule is not None
        self.assertIsNotNone(relaxed_mol)

        # Assert that the relaxed molecule has coordinates
        self.assertTrue(relaxed_mol.GetNumConformers() > 0)


class StrainEnergyCalculationTest(unittest.TestCase):
    def test_get_strain_energy(self):
        # Define a test molecule

        # smile for large molecule
        smiles = "CC(CC1=CC=C(C=C1)C(=O)O)C1=CC=C(C=C1)C(=O)O"

        mol = Chem.MolFromSmiles(smiles)
        # add generate conformer using UFF
        AllChem.EmbedMolecule(mol, randomSeed=0xF00D)

        # Calculate the strain energy
        strain_energy = get_strain_energy(mol)

        # Assert that the calculated strain energy is not NaN
        self.assertFalse(np.isnan(strain_energy))


if __name__ == "__main__":
    unittest.main()

tests/test_strain.py
cch1999/posecheck