From 99f8ca795d8228f703dc873e65879ee83ce2e685 Mon Sep 17 00:00:00 2001 From: zadorozk Date: Wed, 15 May 2024 09:00:03 -0400 Subject: [PATCH 1/4] murcko scaffold split --- src/beignet/splits/__init__.py | 5 + src/beignet/splits/_murcko_scaffold_split.py | 103 ++++++++++++++++++ .../splits/test__murcko_scaffold_split.py | 71 ++++++++++++ 3 files changed, 179 insertions(+) create mode 100644 src/beignet/splits/__init__.py create mode 100644 src/beignet/splits/_murcko_scaffold_split.py create mode 100644 tests/beignet/splits/test__murcko_scaffold_split.py diff --git a/src/beignet/splits/__init__.py b/src/beignet/splits/__init__.py new file mode 100644 index 0000000000..e4e078c316 --- /dev/null +++ b/src/beignet/splits/__init__.py @@ -0,0 +1,5 @@ +from ._murcko_scaffold_split import murcko_scaffold_split + +__all__ = [ + "murcko_scaffold_split", +] diff --git a/src/beignet/splits/_murcko_scaffold_split.py b/src/beignet/splits/_murcko_scaffold_split.py new file mode 100644 index 0000000000..43e9292a98 --- /dev/null +++ b/src/beignet/splits/_murcko_scaffold_split.py @@ -0,0 +1,103 @@ +import math +import random +from collections import defaultdict + +from torch.utils.data import Dataset, Subset + +try: + from rdkit import Chem + from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles + + _RDKit_AVAILABLE = True +except (ImportError, ModuleNotFoundError): + _RDKit_AVAILABLE = False + Chem, MurckoScaffoldSmiles = None, None + + +def murcko_scaffold_split( + dataset: Dataset, + smiles: list[str], + test_size: float | int, + *, + seed: int = 0xDEADBEEF, + shuffle: bool = True, + include_chirality: bool = False, +) -> tuple[Subset, Subset]: + """Split a dataset based on Murcko scaffold splitting based + on provided SMILES strings. + + Note that for datasets that are small or not highly diverse, + the final test set may be smaller than the specified test_size. + + Parameters: + dataset (Dataset): The dataset to split. + smiles (list[str]): A list of SMILES strings. + test_size (float | int): The size of the test set. + If float, should be between 0.0 and 1.0. + If int, should be between 0 and len(smiles). + seed (int): The random seed to use for shuffling. + shuffle (bool): Whether to shuffle the indices. + include_chirality (bool): Whether to include chirality in the scaffold. + + Returns: + tuple[Subset, Subset]: The train and test subsets. + + """ + train_idx, test_idx = _murcko_scaffold_split_indices( + smiles, + test_size, + seed=seed, + shuffle=shuffle, + include_chirality=include_chirality, + ) + return Subset(dataset, train_idx), Subset(dataset, test_idx) + + +def _murcko_scaffold_split_indices( + smiles: list[str], + test_size: float | int, + *, + seed: int = 0xDEADBEEF, + shuffle: bool = True, + include_chirality: bool = False, +) -> tuple[list[int], list[int]]: + """Get train and test indices based on Murcko scaffold splitting.""" + if not _RDKit_AVAILABLE: + raise ImportError( + "This function requires RDKit to be installed (pip install rdkit)" + ) + + if ( + isinstance(test_size, int) and (test_size <= 0 or test_size >= len(smiles)) + ) or (isinstance(test_size, float) and (test_size <= 0 or test_size >= 1)): + raise ValueError( + f"Test_size should be a float in (0, 1) or and int < {len(smiles)}." + ) + + if isinstance(test_size, float): + test_size = math.ceil(len(smiles) * test_size) + + scaffolds = defaultdict(list) + + for ind, s in enumerate(smiles): + mol = Chem.MolFromSmiles(s) + if mol is not None: + scaffold = MurckoScaffoldSmiles(mol=mol, includeChirality=include_chirality) + scaffolds[scaffold].append(ind) + + train_idx = [] + test_idx = [] + + if shuffle: + if seed is not None: + random.Random(seed).shuffle(scaffolds) + else: + random.shuffle(scaffolds) + + for index_list in scaffolds.values(): + if len(test_idx) + len(index_list) <= test_size: + test_idx = [*test_idx, *index_list] + else: + train_idx.extend(index_list) + + return train_idx, test_idx diff --git a/tests/beignet/splits/test__murcko_scaffold_split.py b/tests/beignet/splits/test__murcko_scaffold_split.py new file mode 100644 index 0000000000..d511efac2b --- /dev/null +++ b/tests/beignet/splits/test__murcko_scaffold_split.py @@ -0,0 +1,71 @@ +from importlib.util import find_spec +from unittest.mock import MagicMock, patch + +import pytest +from beignet.splits._murcko_scaffold_split import ( + _murcko_scaffold_split_indices, + murcko_scaffold_split, +) +from torch.utils.data import Dataset, Subset + +_RDKit_AVAILABLE = find_spec("rdkit") is not None + + +@pytest.mark.skipif(not _RDKit_AVAILABLE, reason="RDKit is not available") +@patch("beignet.splits._murcko_scaffold_split._murcko_scaffold_split_indices") +def test_murcko_scaffold_split(mock__murcko_scaffold_split_indices): + mock__murcko_scaffold_split_indices.return_value = ([0], [1]) + + mock_dataset = MagicMock(spec=Dataset) + + train_dataset, test_dataset = murcko_scaffold_split( + dataset=mock_dataset, + smiles=["C", "C"], + test_size=0.5, + shuffle=False, + seed=0, + ) + + assert isinstance(train_dataset, Subset) + assert isinstance(test_dataset, Subset) + assert train_dataset.indices == [0] + assert test_dataset.indices == [1] + + +@pytest.mark.skipif(not _RDKit_AVAILABLE, reason="RDKit is not available") +@pytest.mark.parametrize( + "test_size, expected_train_idx, expected_test_idx", + [ + pytest.param(0.5, [2, 3], [0, 1], id="test_size is float"), + pytest.param(2, [2, 3], [0, 1], id="test_size is int"), + ], +) +def test__murcko_scaffold_split_indices( + test_size, expected_train_idx, expected_test_idx +): + smiles = ["C1CCCCC1", "C1CCCCC1", "CCO", "CCO"] + + train_idx, test_idx = _murcko_scaffold_split_indices( + smiles, + test_size=test_size, + ) + assert train_idx == expected_train_idx + assert test_idx == expected_test_idx + + +@pytest.mark.skipif(not _RDKit_AVAILABLE, reason="RDKit is not available") +@pytest.mark.parametrize( + "smiles, test_size", + [ + pytest.param(["CCO"], 1.2, id="test_size is float > 1"), + pytest.param(["CCO"], -1, id="test_size is negative"), + pytest.param(["CCO"], 0, id="test_size is 0"), + pytest.param(["CCO"], 5, id="test_size > len(smiles)"), + ], +) +def test__murcko_scaffold_split_indices_invalid_inputs(smiles, test_size): + with pytest.raises(ValueError): + _murcko_scaffold_split_indices( + smiles, + test_size=test_size, + ) From d342be429c3ef0824f58574cbf6611ae679d1492 Mon Sep 17 00:00:00 2001 From: zadorozk Date: Wed, 15 May 2024 09:02:49 -0400 Subject: [PATCH 2/4] murcko scaffold split --- src/beignet/splits/_murcko_scaffold_split.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/beignet/splits/_murcko_scaffold_split.py b/src/beignet/splits/_murcko_scaffold_split.py index 43e9292a98..ee8b08cce7 100644 --- a/src/beignet/splits/_murcko_scaffold_split.py +++ b/src/beignet/splits/_murcko_scaffold_split.py @@ -42,6 +42,13 @@ def murcko_scaffold_split( Returns: tuple[Subset, Subset]: The train and test subsets. + References: + - Bemis, G. W., & Murcko, M. A. (1996). + The properties of known drugs. 1. Molecular frameworks. + Journal of medicinal chemistry, 39(15), 2887–2893. + https://doi.org/10.1021/jm9602928 + - "RDKit: Open-source cheminformatics. https://www.rdkit.org" + """ train_idx, test_idx = _murcko_scaffold_split_indices( smiles, From a2cec52a82d59bbb04c7eddb536f9d7a5c3cfd43 Mon Sep 17 00:00:00 2001 From: karinazad Date: Fri, 17 May 2024 16:00:54 -0400 Subject: [PATCH 3/4] doc, subsets --- docs/index.md | 1 + src/beignet/{splits => subsets}/__init__.py | 0 .../_murcko_scaffold_split.py | 62 +++++++++++-------- .../test__murcko_scaffold_split.py | 4 +- 4 files changed, 40 insertions(+), 27 deletions(-) rename src/beignet/{splits => subsets}/__init__.py (100%) rename src/beignet/{splits => subsets}/_murcko_scaffold_split.py (62%) rename tests/beignet/{splits => subsets}/test__murcko_scaffold_split.py (94%) diff --git a/docs/index.md b/docs/index.md index 53f1199f37..947f5afa4a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -42,3 +42,4 @@ ::: beignet.rotation_vector_to_quaternion ::: beignet.rotation_vector_to_rotation_matrix ::: beignet.translation_identity +::: beignet.subsets.murcko_scaffold_split \ No newline at end of file diff --git a/src/beignet/splits/__init__.py b/src/beignet/subsets/__init__.py similarity index 100% rename from src/beignet/splits/__init__.py rename to src/beignet/subsets/__init__.py diff --git a/src/beignet/splits/_murcko_scaffold_split.py b/src/beignet/subsets/_murcko_scaffold_split.py similarity index 62% rename from src/beignet/splits/_murcko_scaffold_split.py rename to src/beignet/subsets/_murcko_scaffold_split.py index ee8b08cce7..8d610d154f 100644 --- a/src/beignet/splits/_murcko_scaffold_split.py +++ b/src/beignet/subsets/_murcko_scaffold_split.py @@ -1,12 +1,13 @@ import math import random from collections import defaultdict +from typing import Sequence from torch.utils.data import Dataset, Subset try: from rdkit import Chem - from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles + from rdkit.Chem.Scaffolds.MurckoScaffold import GetScaffoldForMol _RDKit_AVAILABLE = True except (ImportError, ModuleNotFoundError): @@ -16,39 +17,47 @@ def murcko_scaffold_split( dataset: Dataset, - smiles: list[str], + smiles: Sequence[str], test_size: float | int, *, seed: int = 0xDEADBEEF, shuffle: bool = True, include_chirality: bool = False, ) -> tuple[Subset, Subset]: - """Split a dataset based on Murcko scaffold splitting based + """ + Creates datasets subsets with disjoint Murcko scaffolds based on provided SMILES strings. Note that for datasets that are small or not highly diverse, the final test set may be smaller than the specified test_size. - Parameters: - dataset (Dataset): The dataset to split. - smiles (list[str]): A list of SMILES strings. - test_size (float | int): The size of the test set. - If float, should be between 0.0 and 1.0. - If int, should be between 0 and len(smiles). - seed (int): The random seed to use for shuffling. - shuffle (bool): Whether to shuffle the indices. - include_chirality (bool): Whether to include chirality in the scaffold. - - Returns: - tuple[Subset, Subset]: The train and test subsets. - - References: - - Bemis, G. W., & Murcko, M. A. (1996). - The properties of known drugs. 1. Molecular frameworks. - Journal of medicinal chemistry, 39(15), 2887–2893. - https://doi.org/10.1021/jm9602928 - - "RDKit: Open-source cheminformatics. https://www.rdkit.org" - + Parameters + ---------- + dataset : Dataset + The dataset to split. + smiles : Sequence[str] + A list of SMILES strings. + test_size : float | int + The size of the test set. If float, should be between 0.0 and 1.0. + If int, should be between 0 and len(smiles). + seed : int, optional + The random seed to use for shuffling, by default 0xDEADBEEF + shuffle : bool, optional + Whether to shuffle the indices, by default True + include_chirality : bool, optional + Whether to include chirality in the scaffold, by default False + + Returns + ------- + tuple[Subset, Subset] + The train and test subsets. + + References + ---------- + - Bemis, G. W., & Murcko, M. A. (1996). The properties of known drugs. + 1. Molecular frameworks. Journal of medicinal chemistry, 39(15), 2887–2893. + https://doi.org/10.1021/jm9602928 + - "RDKit: Open-source cheminformatics. https://www.rdkit.org" """ train_idx, test_idx = _murcko_scaffold_split_indices( smiles, @@ -68,7 +77,8 @@ def _murcko_scaffold_split_indices( shuffle: bool = True, include_chirality: bool = False, ) -> tuple[list[int], list[int]]: - """Get train and test indices based on Murcko scaffold splitting.""" + """ + Get train and test indices based on Murcko scaffolds.""" if not _RDKit_AVAILABLE: raise ImportError( "This function requires RDKit to be installed (pip install rdkit)" @@ -89,7 +99,9 @@ def _murcko_scaffold_split_indices( for ind, s in enumerate(smiles): mol = Chem.MolFromSmiles(s) if mol is not None: - scaffold = MurckoScaffoldSmiles(mol=mol, includeChirality=include_chirality) + scaffold = Chem.MolToSmiles( + GetScaffoldForMol(mol), isomericSmiles=include_chirality + ) scaffolds[scaffold].append(ind) train_idx = [] diff --git a/tests/beignet/splits/test__murcko_scaffold_split.py b/tests/beignet/subsets/test__murcko_scaffold_split.py similarity index 94% rename from tests/beignet/splits/test__murcko_scaffold_split.py rename to tests/beignet/subsets/test__murcko_scaffold_split.py index d511efac2b..ab02901a97 100644 --- a/tests/beignet/splits/test__murcko_scaffold_split.py +++ b/tests/beignet/subsets/test__murcko_scaffold_split.py @@ -2,7 +2,7 @@ from unittest.mock import MagicMock, patch import pytest -from beignet.splits._murcko_scaffold_split import ( +from beignet.subsets._murcko_scaffold_split import ( _murcko_scaffold_split_indices, murcko_scaffold_split, ) @@ -12,7 +12,7 @@ @pytest.mark.skipif(not _RDKit_AVAILABLE, reason="RDKit is not available") -@patch("beignet.splits._murcko_scaffold_split._murcko_scaffold_split_indices") +@patch("beignet.subsets._murcko_scaffold_split._murcko_scaffold_split_indices") def test_murcko_scaffold_split(mock__murcko_scaffold_split_indices): mock__murcko_scaffold_split_indices.return_value = ([0], [1]) From d7505d458bb2e220b8dea4f306007f988e9e53f5 Mon Sep 17 00:00:00 2001 From: karinazad Date: Fri, 17 May 2024 16:03:23 -0400 Subject: [PATCH 4/4] docs --- docs/beignet.subsets.md | 3 +++ docs/index.md | 1 - 2 files changed, 3 insertions(+), 1 deletion(-) create mode 100644 docs/beignet.subsets.md diff --git a/docs/beignet.subsets.md b/docs/beignet.subsets.md new file mode 100644 index 0000000000..673ff47254 --- /dev/null +++ b/docs/beignet.subsets.md @@ -0,0 +1,3 @@ +# beignet subsets + +::: beignet.subsets.murcko_scaffold_split \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index 947f5afa4a..53f1199f37 100644 --- a/docs/index.md +++ b/docs/index.md @@ -42,4 +42,3 @@ ::: beignet.rotation_vector_to_quaternion ::: beignet.rotation_vector_to_rotation_matrix ::: beignet.translation_identity -::: beignet.subsets.murcko_scaffold_split \ No newline at end of file