Repository: MachineLearningLifeScience/poli (language: python)

Issue:
Add `dockstring` as an objective function
[Dockstring](https://pubs.acs.org/doi/pdf/10.1021/acs.jcim.1c01334) allows us to compute the docking affinity of small molecules w.r.t. several proteins.


Discussion / hints:

Fix patch:
diff --git a/examples/small_molecules/using_dockstring/querying_dockstring_black_box.py b/examples/small_molecules/using_dockstring/querying_dockstring_black_box.py
new file mode 100644
index 00000000..981c8a3a
--- /dev/null
+++ b/examples/small_molecules/using_dockstring/querying_dockstring_black_box.py
@@ -0,0 +1,15 @@
+"""
+In this small example, we create a black box
+function for dockstring with the drd2 target.
+"""
+
+from poli import objective_factory
+
+if __name__ == "__main__":
+    problem_info, f_dockstring, x0, y0, _ = objective_factory.create(
+        name="dockstring",
+        target_name="DRD2",
+        string_representation="SMILES",
+    )
+
+    print(f"Score of Risperidone: {f_dockstring(x0)}")
diff --git a/src/poli/core/chemistry/tdc_black_box.py b/src/poli/core/chemistry/tdc_black_box.py
index df369595..9c9eac1f 100644
--- a/src/poli/core/chemistry/tdc_black_box.py
+++ b/src/poli/core/chemistry/tdc_black_box.py
@@ -2,7 +2,10 @@
 In this module, we implement a synthetic-accessibility 
 objective using the TDC oracles [1].
 
-[1] TODO: add reference.
+[1] Huang, Kexin, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W. Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. “Artificial Intelligence Foundation for Therapeutic Science.” Nature Chemical Biology 18, no. 10 (October 2022): 1033-36. https://doi.org/10.1038/s41589-022-01131-2.
+
+See also:
+    - The website for TDC: https://tdcommons.ai/
 """
 
 import numpy as np
diff --git a/src/poli/objective_repository/__init__.py b/src/poli/objective_repository/__init__.py
index 17cd6988..315d2abd 100644
--- a/src/poli/objective_repository/__init__.py
+++ b/src/poli/objective_repository/__init__.py
@@ -106,3 +106,10 @@
     AVAILABLE_PROBLEM_FACTORIES["rasp"] = RaspProblemFactory
 except (ImportError, FileNotFoundError):
     pass
+
+try:
+    from .dockstring.register import DockstringProblemFactory
+
+    AVAILABLE_PROBLEM_FACTORIES["dockstring"] = DockstringProblemFactory
+except ImportError:
+    pass
diff --git a/src/poli/objective_repository/dockstring/__init__.py b/src/poli/objective_repository/dockstring/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/poli/objective_repository/dockstring/environment.yml b/src/poli/objective_repository/dockstring/environment.yml
new file mode 100644
index 00000000..5ae9f74e
--- /dev/null
+++ b/src/poli/objective_repository/dockstring/environment.yml
@@ -0,0 +1,13 @@
+name: poli__dockstring
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.9
+  - openbabel
+  - pip
+  - pip:
+    - "git+https://github.com/MachineLearningLifeScience/poli.git@dev"
+    - rdkit
+    - selfies
+    - dockstring
\ No newline at end of file
diff --git a/src/poli/objective_repository/dockstring/register.py b/src/poli/objective_repository/dockstring/register.py
new file mode 100644
index 00000000..e31b87c7
--- /dev/null
+++ b/src/poli/objective_repository/dockstring/register.py
@@ -0,0 +1,139 @@
+"""
+This script implements and registers a black box
+objective function (and a repository) for dockstring [1].
+
+[1] García-Ortegón, Miguel, Gregor N. C. Simm, Austin J. Tripp, José Miguel Hernández-Lobato, Andreas Bender, and Sergio Bacallado. “DOCKSTRING: Easy Molecular Docking Yields Better Benchmarks for Ligand Design.” Journal of Chemical Information and Modeling 62, no. 15 (August 8, 2022): 3486-3502. https://doi.org/10.1021/acs.jcim.1c01334.
+
+"""
+from typing import Tuple
+
+import numpy as np
+
+import selfies as sf
+
+from dockstring import load_target
+
+from poli.core.abstract_black_box import AbstractBlackBox
+from poli.core.abstract_problem_factory import AbstractProblemFactory
+from poli.core.problem_setup_information import ProblemSetupInformation
+
+from poli.core.util.chemistry.string_to_molecule import (
+    translate_selfies_to_smiles,
+    translate_smiles_to_selfies,
+)
+
+from poli.core.util.seeding import seed_numpy, seed_python
+
+
+class DockstringBlackBox(AbstractBlackBox):
+    def __init__(
+        self,
+        info: ProblemSetupInformation,
+        batch_size: int = None,
+        parallelize: bool = False,
+        num_workers: int = None,
+        target_name: str = None,
+        string_representation: str = "SMILES",
+    ):
+        assert (
+            target_name is not None
+        ), "Missing mandatory keyword argument 'target_name'. "
+
+        super().__init__(
+            info=info,
+            batch_size=batch_size,
+            parallelize=parallelize,
+            num_workers=num_workers,
+        )
+        self.target_name = target_name
+        self.string_representation = string_representation
+
+        self.target = load_target(target_name)
+
+    def _black_box(self, x: np.ndarray, context=None) -> np.ndarray:
+        """
+        Assuming that x is an array of strings (either in
+        SMILES or in SELFIES representation).
+        """
+        assert len(x.shape) == 2, "Expected a 2D array of strings. "
+        molecules_as_strings = ["".join(x_i) for x_i in x]
+
+        if self.string_representation == "SELFIES":
+            molecules_as_smiles = translate_selfies_to_smiles(molecules_as_strings)
+        else:
+            molecules_as_smiles = molecules_as_strings
+
+        # TODO: Should we parallelize?
+        scores = [self.target.dock(smiles)[0] for smiles in molecules_as_smiles]
+
+        # Since our goal is maximization, and scores in dockstring
+        # are better if they are lower, we return the negative of
+        # the scores.
+        return -np.array(scores).reshape(-1, 1)
+
+
+class DockstringProblemFactory(AbstractProblemFactory):
+    def get_setup_information(self) -> ProblemSetupInformation:
+        # TODO: We might change this in the future for a
+        # default dictionary, depending on whether we
+        # are using SMILES or SELFIES.
+        alphabet = None
+
+        return ProblemSetupInformation(
+            name="dockstring",
+            max_sequence_length=np.inf,
+            aligned=False,
+            alphabet=alphabet,
+        )
+
+    def create(
+        self,
+        seed: int = None,
+        batch_size: int = None,
+        parallelize: bool = False,
+        num_workers: int = None,
+        target_name: str = None,
+        string_representation: str = "SMILES",
+    ) -> Tuple[AbstractBlackBox, np.ndarray, np.ndarray]:
+        assert (
+            target_name is not None
+        ), "Missing mandatory keyword argument 'target_name'. "
+
+        seed_numpy(seed)
+        seed_python(seed)
+
+        dockstring_black_box = DockstringBlackBox(
+            info=self.get_setup_information(),
+            batch_size=batch_size,
+            parallelize=parallelize,
+            num_workers=num_workers,
+            target_name=target_name,
+            string_representation=string_representation,
+        )
+
+        # Using the initial example they provide in the
+        # dockstring documentation (i.e. DRD2 and Risperidone)
+        risperidone_smiles = "CC1=C(C(=O)N2CCCCC2=N1)CCN3CCC(CC3)C4=NOC5=C4C=CC(=C5)F"
+        if string_representation.upper() == "SMILES":
+            # TODO: replace for proper smiles tokenization.
+            x0 = np.array([list(risperidone_smiles)])
+        elif string_representation.upper() == "SELFIES":
+            risperidone_selfies = translate_smiles_to_selfies([risperidone_smiles])[0]
+            risperidone_selfies_as_tokens = list(sf.split_selfies(risperidone_selfies))
+            x0 = np.array([risperidone_selfies_as_tokens])
+        else:
+            raise ValueError(
+                f"Invalid string representation. Expected SMILES or SELFIES but received {string_representation}."
+            )
+
+        return dockstring_black_box, x0, dockstring_black_box(x0)
+
+
+if __name__ == "__main__":
+    from poli.core.registry import register_problem
+
+    dockstring_problem_factory = DockstringProblemFactory()
+    register_problem(
+        dockstring_problem_factory,
+        conda_environment_name="poli__dockstring",
+    )
diff --git a/tox.dev.ini b/tox.dev.ini
index 37a38b4c..c5cb0e13 100644
--- a/tox.dev.ini
+++ b/tox.dev.ini
@@ -21,6 +21,9 @@ commands =
     sh -c 'if conda info --envs | grep -q poli__protein; then echo "poli__protein already exists"; else conda env create -f ./src/poli/objective_repository/foldx_stability/environment.yml; fi'
     sh -c "conda run -n poli__protein python -m pip uninstall -y poli"
     sh -c "conda run -n poli__protein python -m pip install -e ."
+    sh -c 'if conda info --envs | grep -q poli__dockstring; then echo "poli__protein already exists"; else conda env create -f ./src/poli/objective_repository/dockstring/environment.yml; fi'
+    sh -c "conda run -n poli__dockstring python -m pip uninstall -y poli"
+    sh -c "conda run -n poli__dockstring python -m pip install -e ."
     pytest {tty:--color=yes} -v {posargs}
     sh -c "rm -rf ~/.poli_objectives/*.sh"
     sh -c "rm -rf ~/.poli_objectives/config.rc"


Test patch:
diff --git a/src/poli/tests/registry/chemistry/test_chemistry_objectives.py b/src/poli/tests/registry/chemistry/test_chemistry_objectives.py
index 01443904..7ba28017 100644
--- a/src/poli/tests/registry/chemistry/test_chemistry_objectives.py
+++ b/src/poli/tests/registry/chemistry/test_chemistry_objectives.py
@@ -146,3 +146,72 @@ def test_penalized_logp_lambo():
     print(x0)
     print(y0)
     f.terminate()
+
+
+def test_querying_dockstring_using_smiles():
+    """
+    In this test, we force-register and query dockstring.
+    """
+    from poli import objective_factory
+
+    _, f, x0, y0, _ = objective_factory.create(
+        name="dockstring",
+        target_name="DRD2",
+        string_representation="SMILES",
+        force_register=True,
+    )
+
+    # Docking another smiles
+    x1 = np.array([list("CC(=O)OC1=CC=CC=C1C(=O)O")])
+    y1 = f(x1)
+
+    f.terminate()
+
+
+def test_querying_dockstring_using_selfies():
+    """
+    In this test, we check whether dockstring still
+    works when using SELFIES instead of SMILES.
+    """
+    from poli import objective_factory
+
+    _, f, x0, y0, _ = objective_factory.create(
+        name="dockstring",
+        target_name="ABL1",
+        string_representation="SELFIES",
+        force_register=True,
+    )
+
+    # Docking another smiles
+    selfies_aspirin = np.array(
+        [
+            [
+                "[C]",
+                "[C]",
+                "[=Branch1]",
+                "[C]",
+                "[=O]",
+                "[O]",
+                "[C]",
+                "[=C]",
+                "[C]",
+                "[=C]",
+                "[C]",
+                "[=C]",
+                "[Ring1]",
+                "[=Branch1]",
+                "[C]",
+                "[=Branch1]",
+                "[C]",
+                "[=O]",
+                "[O]",
+            ]
+        ]
+    )
+
+    y1 = f(selfies_aspirin)
+    f.terminate()
+
+
+if __name__ == "__main__":
+    test_querying_dockstring_using_selfies()