Repository: dockstring/dockstring (language: python)

Issue:
Add API for molecular optimization baseline.
Adds:
- `dockstring.benchmarks` module to hold benchmark-related code.
- Implementations of 3 benchmark objectives from original paper
- Testing of the above implementations.


Discussion / hints:

Fix patch:
diff --git a/.mypy.ini b/.mypy.ini
index 43a0a2d..c00c48b 100644
--- a/.mypy.ini
+++ b/.mypy.ini
@@ -18,3 +18,9 @@ ignore_missing_imports = True
 
 [mypy-pandas.*]
 ignore_missing_imports = True
+
+[mypy-pytest.*]
+ignore_missing_imports = True
+
+[mypy-pkg_resources.*]
+ignore_missing_imports = True
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 8aa980d..979ee5c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,11 +1,11 @@
 repos:
 
-# yapf
+# yapf (code formatting)
 -   repo: https://github.com/pre-commit/mirrors-yapf
-    rev: 'v0.31.0'  # Use the sha / tag you want to point at
+    rev: 'v0.32.0'  # Use the sha / tag you want to point at
     hooks:
     -   id: yapf
-        args: [--style=.style.yapf, --recursive, --diff, .]
+        args: [--style=.style.yapf, --recursive, --in-place, .]
 
 # flake8
 -   repo: https://github.com/pycqa/flake8
@@ -14,12 +14,9 @@ repos:
     -   id: flake8
         args: [--config=.flake8]
 
-# mypy
-# For some reason, pre-commit mypy doesn't like our folder structure raising errors such as:
-# Duplicate module named 'dockstring'
-# Disabling for now.
-#-   repo: https://github.com/pre-commit/mirrors-mypy
-#    rev: 'v0.812'  # pick a git hash / tag to point to
-#    hooks:
-#    -   id: mypy
-#        args: [--config=.mypy.ini, ./dockstring]
+# mypy (type checking)
+-  repo: https://github.com/pre-commit/mirrors-mypy
+   rev: 'v0.971'  # pick a git hash / tag to point to
+   hooks:
+   -   id: mypy
+       args: [--config=.mypy.ini]
diff --git a/README.md b/README.md
index e5fbc2b..8982f51 100644
--- a/README.md
+++ b/README.md
@@ -3,23 +3,43 @@
 ![CI Tests](https://github.com/mgarort/dockstring/workflows/Install%20conda%20env%20and%20run%20pytest./badge.svg?branch=main)
 ![Code Style: yapf](https://img.shields.io/badge/code%20style-yapf-orange.svg)
 
-A Python package for easy molecular docking.
+A Python package for easy molecular docking and docking benchmarking.
+We can dock molecules in a few lines of code from just a SMILES string!
+For details, see [our paper](https://pubs.acs.org/doi/full/10.1021/acs.jcim.1c01334)
+and our [website](https://dockstring.github.io/):
+
+> García-Ortegón, Miguel, et al. "DOCKSTRING: easy molecular docking yields better benchmarks for ligand design." Journal of Chemical Information and Modeling (2021).
 
 ## Installation
 
-1. Create a new `conda` environment from the `environment.yml` file in this repository:
-   ```bash
-   conda env create -f environment.yml
-   ```
-1. Activate the new `conda` environment:
+To ensure compatibility with the dockstring dataset,
+the package has very strict versioning requirements
+for its main dependencies (`rdkit` and `openbabel`).
+As such, we recommend you install it in the following way.
+
+1. Clone this repository.
+1. Choose whether to install into an existing environment or create a new environment.
+    - To install into an existing environment, install the correct version of `openbabel` and `rdkit`:
+      ```bash
+      conda install -c conda-forge rdkit=2021.03.3 openbabel=3.1.1
+      ```
+    - To install into a new environment, run:
+      ```bash
+      conda env create -f environment.yml
+      conda activate dockstring
+      ```
+1. Install the dockstring package with `pip` from this repository:
    ```bash
-   conda activate dockstring
+   pip install .
    ```
-1. Install the dockstring package with `pip`:
+1. Check whether the installation was successful by running a test script.
+   Running without error inducates a successful install.
    ```bash
-   pip install .
+   python tutorials/simple_example.py
    ```
 
+If this method of installation does not work for you, please raise a github issue and we will try to help.
+
 ### Optional
 
 Install [PyMol](https://pymol.org/) for target, search box and ligand visualization:
@@ -27,27 +47,53 @@ Install [PyMol](https://pymol.org/) for target, search box and ligand visualizat
 conda install -c conda-forge pymol-open-source 
 ```
 
-## Tutorial
+## Tutorials
+
+- See dockstring's basic usage [here](tutorials/1_docking_risperidone_against_DRD2.ipynb).
+- Learn how to visualize docking poses [here](tutorials/2_visualizing_dataset_poses.ipynb)
 
-Check out this [tutorial](tutorials/1_docking_risperidone_against_DRD2.ipynb) to get started!
+See [our website](https://dockstring.github.io/) for linkks to tutorials for
+our dataset and benchmarks.
 
 ## Development
 
-### Code format
+We use [pre-commit](https://pre-commit.com/) to enforce code formatting and style.
+Install by running:
 
-We use yapf and flake8 for code formatting.
-Run the following to check formatting:
+```bash
+conda install -c conda-forge pre-commit
+pre-commit install
+```
+
+We use [pytest](https://docs.pytest.org) to test our code.
+You can install pytest by running `conda install -c conda-forge pytest`.
+Before committing, please run the following to make sure that all tests pass:
 
 ```bash
-yapf --style=.style.yapf --in-place --recursive .
-flake8 --config=.flake8 .
+python -m pytest tests/
 ```
 
-We have CI set up to check this, but we _highly_ recommend setting up
-[pre-commit](https://pre-commit.com/) to avoid accidentally committing bad code.
-You can do so in the following way:
+Alternatively, to skip a variety of slow tests, run:
 
 ```bash
-conda install -c conda-forge pre-commit
-pre-commit install
+python -m pytest -m "not slow" tests/
+```
+
+## Citation
+
+If you use the dockstring package/dataset/benchmark in your work,
+please use the following citation:
+
+```tex
+@article{garciaortegon2022dockstring,
+    author = {García-Ortegón, Miguel and Simm, Gregor N. C. and Tripp, Austin J. and Hernández-Lobato, José Miguel and Bender, Andreas and Bacallado, Sergio},
+    title = {DOCKSTRING: Easy Molecular Docking Yields Better Benchmarks for Ligand Design},
+    journal = {Journal of Chemical Information and Modeling},
+    volume = {62},
+    number = {15},
+    pages = {3486-3502},
+    year = {2022},
+    doi = {10.1021/acs.jcim.1c01334},
+    URL = {https://doi.org/10.1021/acs.jcim.1c01334}
+}
 ```
diff --git a/dockstring/_version.py b/dockstring/_version.py
index b794fd4..7fd229a 100644
--- a/dockstring/_version.py
+++ b/dockstring/_version.py
@@ -1,1 +1,1 @@
-__version__ = '0.1.0'
+__version__ = '0.2.0'
diff --git a/dockstring/benchmarks/__init__.py b/dockstring/benchmarks/__init__.py
new file mode 100644
index 0000000..786898a
--- /dev/null
+++ b/dockstring/benchmarks/__init__.py
@@ -0,0 +1,3 @@
+from . import original, utils
+
+__all__ = ["original", "utils"]
diff --git a/dockstring/benchmarks/original.py b/dockstring/benchmarks/original.py
new file mode 100644
index 0000000..435ad81
--- /dev/null
+++ b/dockstring/benchmarks/original.py
@@ -0,0 +1,79 @@
+"""Benchmarks for original dockstring paper."""
+
+import functools
+import math
+from typing import Dict
+
+from .utils import BenchmarkObjective, safe_dock_function, QED as QED_function
+
+
+# Raw functions
+def QED_penalty(qed: float) -> float:
+    return 10.0 * (1.0 - qed)
+
+
+def F2_score(*, F2: float, QED: float) -> float:
+    return F2 + QED_penalty(QED)
+
+
+def promiscuous_PPAR_score(*, PPARA: float, PPARD: float, PPARG: float, QED: float) -> float:
+    # Max of a list of NaNs is not always NaN so we check this manually
+    if any(math.isnan(v) for v in [PPARA, PPARD, PPARG]):
+        return math.nan
+    return max(PPARA, PPARD, PPARG) + QED_penalty(QED)
+
+
+def selective_JAK2_score(*, JAK2: float, LCK: float, QED: float) -> float:
+
+    # Note: there was a small error in the formula for this objective
+    # in our JCIM publication. *THIS* is the correct formula,
+    # which matches the numbers in all of our tables
+    lck_median_score = -8.1
+    return JAK2 - min(LCK - lck_median_score, 0) + QED_penalty(QED)
+
+
+def get_benchmark_functions(**dock_kwargs) -> Dict[str, BenchmarkObjective]:
+    """
+    Returns the functions for the original benckmarks.
+
+    dock_kwargs specifies kwargs to pass to the `target.dock` function
+    (e.g. pH, num_cpus)
+    """
+    output: Dict[str, BenchmarkObjective] = dict()
+
+    # F2
+    output["F2"] = BenchmarkObjective(
+        base_functions=dict(
+            F2=functools.partial(safe_dock_function, target_name="F2", **dock_kwargs),
+            QED=QED_function,
+        ),
+        aggregation_function=F2_score,
+    )
+
+    # Promiscuous PPAR
+    ppar_funcs = {
+        target_name: functools.partial(safe_dock_function, target_name=target_name, **dock_kwargs)
+        for target_name in ["PPARA", "PPARD", "PPARG"]
+    }
+    output["promiscuous_PPAR"] = BenchmarkObjective(
+        base_functions=dict(
+            QED=QED_function,
+            **ppar_funcs,
+        ),
+        aggregation_function=promiscuous_PPAR_score,
+    )
+
+    # Selective JAK2
+    jak2_funcs = {
+        target_name: functools.partial(safe_dock_function, target_name=target_name, **dock_kwargs)
+        for target_name in ["JAK2", "LCK"]
+    }
+    output["selective_JAK2"] = BenchmarkObjective(
+        base_functions=dict(
+            QED=QED_function,
+            **jak2_funcs,
+        ),
+        aggregation_function=selective_JAK2_score,
+    )
+
+    return output
diff --git a/dockstring/benchmarks/utils.py b/dockstring/benchmarks/utils.py
new file mode 100644
index 0000000..5be173d
--- /dev/null
+++ b/dockstring/benchmarks/utils.py
@@ -0,0 +1,45 @@
+"""Some generally useful functions."""
+
+from dataclasses import dataclass
+from typing import Callable, Dict, Tuple
+
+from rdkit import Chem
+from rdkit.Chem import QED as qed_module
+
+import dockstring
+
+
+@dataclass
+class BenchmarkObjective:
+    """
+    General class for an objective function whose calculation requires
+    1) evaluating a set of independent base functions
+    2) aggregating those scores together into a single score
+    """
+    base_functions: Dict[str, Callable[[str], float]]
+    aggregation_function: Callable[..., float]
+
+    def _eval_base_functions(self, smiles: str) -> Dict[str, float]:
+        return {name: f(smiles) for name, f in self.base_functions.items()}
+
+    def __call__(self, smiles: str) -> Tuple[float, Dict[str, float]]:
+        """Call all in 1."""
+        base_fn_vals = self._eval_base_functions(smiles)
+        return self.aggregation_function(**base_fn_vals), base_fn_vals
+
+
+def safe_dock_function(smiles: str, target_name: str, **dock_kwargs):
+    """Call dockstring and return nan if there are errors."""
+    target = dockstring.load_target(target_name)
+    try:
+        docking_output = target.dock(smiles, **dock_kwargs)
+        score = docking_output[0]
+    except dockstring.DockstringError:
+        score = float("nan")
+    return score
+
+
+def QED(smiles: str) -> float:
+    """Calculates QED from a SMILES string."""
+    mol = Chem.MolFromSmiles(smiles)
+    return qed_module.qed(mol)
diff --git a/tutorials/simple_example.py b/tutorials/simple_example.py
new file mode 100644
index 0000000..e374d2c
--- /dev/null
+++ b/tutorials/simple_example.py
@@ -0,0 +1,8 @@
+"""Simple example of invoking dockstring."""
+
+from dockstring import load_target
+
+if __name__ == "__main__":
+    target = load_target("DRD2")
+    score, _ = target.dock("CC1=C(C(=O)N2CCCCC2=N1)CCN3CCC(CC3)C4=NOC5=C4C=CC(=C5)F")
+    print(f"Docking was successful, score={score:.3g}. END OF SCRIPT.")


Test patch:
diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml
index 8af79ce..ed68144 100644
--- a/.github/workflows/run-tests.yml
+++ b/.github/workflows/run-tests.yml
@@ -53,4 +53,4 @@ jobs:
       - name: Test with pytest
         run: |
           conda activate dockstring
-          python3 -m pytest tests
+          python3 -m pytest -m "not slow" tests
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..a83f962
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,3 @@
+[pytest]
+markers =
+    slow: mark test as slow (too slow for CLI)
\ No newline at end of file
diff --git a/test-environment.yml b/test-environment.yml
index f4e5497..8dd1b4d 100644
--- a/test-environment.yml
+++ b/test-environment.yml
@@ -5,4 +5,4 @@ dependencies:
   - rdkit=2021.03.3
   - openbabel=3.1.1
   - python=3.7.10
-  - pytest
+  - pytest
\ No newline at end of file
diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py
new file mode 100644
index 0000000..49906e6
--- /dev/null
+++ b/tests/test_benchmarks.py
@@ -0,0 +1,80 @@
+"""Simple tests for the benchmarks."""
+
+import functools
+import math
+from typing import Dict
+
+import pytest
+
+import dockstring.benchmarks
+
+
+@pytest.fixture
+def original_benchmarks():
+    return dockstring.benchmarks.original.get_benchmark_functions()
+
+
+# For convenience, define a function for "good enough" float comparison
+close_enough = functools.partial(math.isclose, rel_tol=1e-4)
+
+
+class TestOriginalBenchmarks:
+    """Tests for the original benchmarks."""
+    def test_smoke(self, original_benchmarks: dict):
+        """Can I load the benchmark functions?"""
+        assert set(original_benchmarks.keys()) == {"F2", "promiscuous_PPAR", "selective_JAK2"}
+
+    @pytest.mark.slow
+    @pytest.mark.parametrize("name,top_smiles,expected_obj_parts,expected_obj", [
+        ("F2", "CC1=CC2(c3cccc(C(=O)NC4=NC(=O)CC(F)=N4)c3)CC(=C1)C=CC2=O", dict(F2=-12.7, QED=0.85955), -11.2955),
+        ("promiscuous_PPAR", "O=C1C=CC2=CC=C3C=C(Nc4cccc(C5C=c6cc7c(cc6=C5)CC=C7)c4)N=C3C12",
+         dict(PPARA=-13.3, PPARD=-13.8, PPARG=-13.3, QED=0.829705), -11.59705),
+        (
+            "selective_JAK2",
+            "CC1=NCC(C(N)=O)=NN1c1cc(-c2c(C)cc(C)cc2C)ccc1F",
+            dict(JAK2=-12.2, LCK=-8.3, QED=0.9182718),
+            -11.182718,
+        ),
+    ])
+    def test_benchmark_values(self, original_benchmarks: dict, name: str, top_smiles: str,
+                              expected_obj_parts: Dict[str, float], expected_obj: float):
+        """
+        Test the values of all the benchmark functions using top-scoring SMILES from the paper.
+        """
+
+        benchmark_fn = original_benchmarks[name]
+
+        # First, check that the score formula seems correct
+        # (i.e. if the parts of the score are correct, do they form the correct whole)
+        assert close_enough(expected_obj, benchmark_fn.aggregation_function(**expected_obj_parts))
+
+        # Second, check end-to end whether the SMILES produces the score expected
+        actual_obj, actual_obj_parts = benchmark_fn(top_smiles)
+        assert close_enough(actual_obj, expected_obj)  # overall value matches
+        # Ensure parts are the same up to a certain tolerance
+        assert set(actual_obj_parts.keys()) == set(expected_obj_parts.keys())
+        for key in actual_obj_parts.keys():
+            assert close_enough(actual_obj_parts[key], expected_obj_parts[key])
+
+    @pytest.mark.slow
+    @pytest.mark.parametrize(
+        "name,nan_smiles",
+        [
+            ("F2", "S=C1N(C(O)C(=C(C2C(=O)N(C(=S)N(C2=O)C)C)C=3C=CN=CC3)C(=O)N1C)C"),
+            ("promiscuous_PPAR", "OC1=C(C(N[C@H]2[C@H](N=C(C3=C(O)C=CC=C3)C)CCCC2)=C)C=CC=C1"),  # not NaN for first one
+            ("promiscuous_PPAR", "O=C1C2(CN3C4(N(C2)CC1(C3)C(C)C)C5=C(NC4=O)C=CC=C5)C(C)C"),  # nan for all
+            ("selective_JAK2", "S=C1N(C(O)C(=C(C2C(=O)N(C(=S)N(C2=O)C)C)C=3C=CN=CC3)C(=O)N1C)C"),  # nan for JAK2 only
+            ("selective_JAK2", "O=C1C2(CN3C4(N(C2)CC1(C3)C(C)C)C5=C(NC4=O)C=CC=C5)C(C)C"),  # nan for LCK only
+        ])
+    def test_benchmark_nan(self, original_benchmarks: dict, name: str, nan_smiles: str):
+        """
+        Make sure everything is ok for SMILES with NaN docking scores.
+        The NaN smiles are sourced from the dockstring dataset itself.
+        """
+        benchmark_fn = original_benchmarks[name]
+
+        # Test 1: does it run ok?
+        value, _ = benchmark_fn(nan_smiles)
+
+        # Test 2: is the value actually nan?
+        assert math.isnan(value)
diff --git a/tests/test_targets.py b/tests/test_targets.py
index 78d5f95..17b46c0 100644
--- a/tests/test_targets.py
+++ b/tests/test_targets.py
@@ -139,17 +139,22 @@ def test_kekulization(self, smiles):
         assert refine_mol_with_ff(embedded_mol)
 
 
-class TestDocking:
-    def test_simple_docking(self):
-        target = load_target('ABL1')
+def test_simple_docking():
+    """Simple docking test, set to run in all cases."""
+    target = load_target('ABL1')
+
+    smiles_1 = 'CCO'
+    energy_1, _ = target.dock(smiles_1)
+    assert energy_1 is not None and math.isclose(energy_1, -2.4)
+
+    smiles_2 = 'CC'
+    energy_2, _ = target.dock(smiles_2)
+    assert energy_2 is not None and math.isclose(energy_2, -1.8)
 
-        smiles_1 = 'CCO'
-        energy_1, _ = target.dock(smiles_1)
-        assert energy_1 is not None and math.isclose(energy_1, -2.4)
 
-        smiles_2 = 'CC'
-        energy_2, _ = target.dock(smiles_2)
-        assert energy_2 is not None and math.isclose(energy_2, -1.8)
+@pytest.mark.slow
+class TestDocking_Slow:
+    """Run a large number of slower docking tests."""
 
     # Test different SMILES representations of lysine
     @pytest.mark.parametrize('smiles', [lysine_smiles, 'NCCCC[C@H](N)C(=O)O'])
@@ -271,24 +276,11 @@ def test_bond_assignment_fails(self, target_name: str, ligand_smiles: str):
         with pytest.raises(PoseProcessingError):
             target.dock(ligand_smiles)
 
-    # Commented out because takes too long
-    # @pytest.mark.parametrize('target_name, ligand', [
-    #     ('ABL1', 'S(=O)(=O)(N(C[C@@H]1OCCCC[C@@H](OC=2C(C(=O)N(C[C@H]1C)[C@@H](CO)C)=CC(NC(=O)NC=3C=CC(F)=CC3)=CC2)C)C)C=4SC=CC4'),
-    # ])
-    # def test_atom_valence_exception(self, target_name: str, ligand: str):
-    #     target = load_target(target_name, working_dir='/home/gregor/.config/JetBrains/PyCharm2021.1/scratches/failure')
-    #     assert target.dock(ligand)
-
-    # Commented out because test takes too long
-    """
-    @pytest.mark.parametrize(
-        'target, ligand',
-        [
-            ('PTPN1', 'O=S(=O)(C1=CC=C(N(C(C2=CC=C(C=C2)NC(C3=CC=C(C=C3)C#N)=O)=O)C)C=C1)NC4=NC=CS4'),
-        ],
-    )
-    def test_docking_failure(self, target, ligand):
-        target = load_target(target)
-        with pytest.raises(DockingError):
-            target.dock(ligand)
-    """
+    @pytest.mark.parametrize('target_name, ligand', [
+        ('ABL1',
+         'S(=O)(=O)(N(C[C@@H]1OCCCC[C@@H](OC=2C(C(=O)N(C[C@H]1C)[C@@H](CO)C)=CC(NC(=O)NC=3C=CC(F)=CC3)=CC2)C)C)C=4SC=CC4'
+         ),
+    ])
+    def test_atom_valence_exception(self, target_name: str, ligand: str):
+        target = load_target(target_name)
+        assert target.dock(ligand)