diff --git a/.github/workflows/docking.yaml b/.github/workflows/docking.yaml index 530ae995..8052fc3c 100644 --- a/.github/workflows/docking.yaml +++ b/.github/workflows/docking.yaml @@ -83,11 +83,8 @@ jobs: - name: Install packages run: | python -m pip install -e ./asapdiscovery-data --no-deps - python -m pip install -e ./asapdiscovery-dataviz --no-deps python -m pip install -e ./asapdiscovery-docking --no-deps - python -m pip install -e ./asapdiscovery-ml --no-deps python -m pip install -e ./asapdiscovery-modeling --no-deps - python -m pip install -e ./asapdiscovery-spectrum --no-deps micromamba list - name: Run tests diff --git a/asapdiscovery-dataviz/asapdiscovery/dataviz/fint_score.py b/asapdiscovery-dataviz/asapdiscovery/dataviz/fint_score.py new file mode 100644 index 00000000..b5702bbb --- /dev/null +++ b/asapdiscovery-dataviz/asapdiscovery/dataviz/fint_score.py @@ -0,0 +1,116 @@ +from pathlib import Path +from typing import ClassVar, Union + +from multimethod import multimethod +from pydantic.v1 import Field, validator + +from asapdiscovery.dataviz.plip import compute_fint_score +from asapdiscovery.docking.docking import DockingResult +from asapdiscovery.docking.scorer import ( + ScorerBase, + ScoreType, + ScoreUnits, + Score, + _get_disk_path_from_docking_result, +) +from asapdiscovery.spectrum.fitness import target_has_fitness_data +from asapdiscovery.data.schema.complex import Complex +from asapdiscovery.data.services.postera.manifold_data_validation import TargetTags +from asapdiscovery.data.util.dask_utils import dask_vmap, backend_wrapper + + +class FINTScorer(ScorerBase): + """ + Score using Fitness Interaction Score + + Overloaded to accept DockingResults, Complexes, or Paths to PDB files. + """ + + score_type: ScoreType = Field(ScoreType.FINT, description="Type of score") + units: ClassVar[ScoreUnits.arbitrary] = ScoreUnits.arbitrary + target: TargetTags = Field(..., description="Which target to use for scoring") + + @validator("target") + @classmethod + def validate_target(cls, v): + if not target_has_fitness_data(v): + raise ValueError( + "target does not have fitness data so cannot use FINTScorer" + ) + return v + + @dask_vmap(["inputs"]) + @backend_wrapper("inputs") + def _score( + self, + inputs: Union[list[DockingResult], list[Complex], list[Path]], + return_for_disk_backend: bool = False, + **kwargs, + ) -> list[Score]: + """ + Score the inputs, dispatching based on type. + """ + return self._dispatch( + inputs, return_for_disk_backend=return_for_disk_backend, **kwargs + ) + + @multimethod + def _dispatch( + self, + inputs: list[DockingResult], + return_for_disk_backend: bool = False, + **kwargs, + ) -> list[Score]: + """ + Dispatch for DockingResults + """ + results = [] + for inp in inputs: + _, fint_score = compute_fint_score( + inp.to_protein(), inp.posed_ligand.to_oemol(), self.target + ) + + sc = Score.from_score_and_docking_result( + fint_score, self.score_type, self.units, inp + ) + # overwrite the input with the path to the file + if return_for_disk_backend: + sc.input = _get_disk_path_from_docking_result(inp) + + results.append(sc) + + return results + + @_dispatch.register + def _dispatch(self, inputs: list[Complex], **kwargs): + """ + Dispatch for Complexes + """ + results = [] + for inp in inputs: + _, fint_score = compute_fint_score( + inp.target.to_oemol(), inp.ligand.to_oemol(), self.target + ) + results.append( + Score.from_score_and_complex( + fint_score, self.score_type, self.units, inp + ) + ) + return results + + @_dispatch.register + def _dispatch(self, inputs: list[Path], **kwargs): + """ + Dispatch for PDB files from disk + """ + # assuming reading PDB files from disk + complexes = [ + Complex.from_pdb( + p, + ligand_kwargs={"compound_name": f"{p.stem}_ligand"}, + target_kwargs={"target_name": f"{p.stem}_target"}, + ) + for p in inputs + ] + + return self._dispatch(complexes, **kwargs) diff --git a/asapdiscovery-dataviz/asapdiscovery/dataviz/tests/test_fintscore.py b/asapdiscovery-dataviz/asapdiscovery/dataviz/tests/test_fintscore.py index be84a586..0fbeea62 100644 --- a/asapdiscovery-dataviz/asapdiscovery/dataviz/tests/test_fintscore.py +++ b/asapdiscovery-dataviz/asapdiscovery/dataviz/tests/test_fintscore.py @@ -1,9 +1,12 @@ from pathlib import Path - +import pytest from asapdiscovery.data.backend.openeye import load_openeye_pdb +from asapdiscovery.dataviz.fint_score import FINTScorer from asapdiscovery.dataviz.plip import compute_fint_score from asapdiscovery.data.readers.molfile import MolFileFactory from asapdiscovery.data.testing.test_resources import fetch_test_file +from asapdiscovery.docking.openeye import POSITDockingResults +from asapdiscovery.data.schema.complex import Complex def test_fint_score(): @@ -32,3 +35,43 @@ def test_fint_score(): # should both fall between 0 and 1 assert 0 <= fint_score[0] <= 1.0 assert 0 <= fint_score[1] <= 1.0 + + +@pytest.fixture() +def results_simple(): + return [ + POSITDockingResults.from_json_file( + fetch_test_file("docking_results_simple.json") + ) + ] + + +@pytest.fixture() +def results_simple_nolist(results_simple): + return results_simple[0] + + +@pytest.fixture() +def complex_simple(): + return Complex.from_pdb( + fetch_test_file("Mpro-P0008_0A_ERI-UCB-ce40166b-17_prepped_receptor_0.pdb"), + ligand_kwargs={"compound_name": "test"}, + target_kwargs={"target_name": "test", "target_hash": "mock_hash"}, + ) + + +@pytest.fixture() +def pdb_simple(): + return fetch_test_file("Mpro-P0008_0A_ERI-UCB-ce40166b-17_prepped_receptor_0.pdb") + + +@pytest.mark.parametrize( + "data_fixture", ["results_simple_nolist", "complex_simple", "pdb_simple"] +) +@pytest.mark.parametrize("return_df", [True, False]) +@pytest.mark.parametrize("use_dask", [True, False]) +def test_FINT_scorer(use_dask, return_df, data_fixture, request): + data = request.getfixturevalue(data_fixture) + scorer = FINTScorer(target="SARS-CoV-2-Mpro") + scores = scorer.score([data], use_dask=use_dask, return_df=return_df) + assert len(scores) == 1 diff --git a/asapdiscovery-docking/asapdiscovery/docking/scorer.py b/asapdiscovery-docking/asapdiscovery/docking/scorer.py index dc60b2dc..be64be03 100644 --- a/asapdiscovery-docking/asapdiscovery/docking/scorer.py +++ b/asapdiscovery-docking/asapdiscovery/docking/scorer.py @@ -10,11 +10,9 @@ import numpy as np import pandas as pd from asapdiscovery.data.backend.openeye import oedocking, oemol_to_pdb_string -from asapdiscovery.dataviz.plip import compute_fint_score from asapdiscovery.data.schema.complex import Complex from asapdiscovery.data.schema.ligand import Ligand, LigandIdentifiers from asapdiscovery.data.schema.target import TargetIdentifiers -from asapdiscovery.data.services.postera.manifold_data_validation import TargetTags from asapdiscovery.data.util.dask_utils import ( BackendType, FailureMode, @@ -23,12 +21,8 @@ ) from asapdiscovery.docking.docking import DockingResult from asapdiscovery.docking.docking_data_validation import DockingResultCols -from asapdiscovery.ml.inference import InferenceBase, get_inference_cls_from_model_type -from asapdiscovery.ml.models import MLModelSpecBase -from asapdiscovery.spectrum.fitness import target_has_fitness_data -from mtenn.config import ModelType from multimethod import multimethod -from pydantic.v1 import BaseModel, Field, validator +from pydantic.v1 import BaseModel, Field logger = logging.getLogger(__name__) @@ -59,42 +53,6 @@ class ScoreUnits(str, Enum): INVALID = "INVALID" -# TODO: this is a massive kludge, need to refactor -def endpoint_and_model_type_to_score_type(endpoint: str, model_type: str) -> ScoreType: - """ - Convert an endpoint to a score type. - - Parameters - ---------- - endpoint : str - Endpoint to convert - - Returns - ------- - ScoreType - Score type - """ - if model_type == ModelType.GAT: - if endpoint == "pIC50": # TODO: make this an enum - return ScoreType.GAT_pIC50 - elif endpoint == "LogD": - return ScoreType.GAT_LogD - else: - raise ValueError(f"Endpoint {endpoint} not recognized, for GAT") - elif model_type == ModelType.schnet: - if endpoint == "pIC50": - return ScoreType.schnet_pIC50 - else: - raise ValueError(f"Endpoint {endpoint} not recognized, for Schnet") - elif model_type == ModelType.e3nn: - if endpoint == "pIC50": - return ScoreType.e3nn_pIC50 - else: - raise ValueError(f"Endpoint {endpoint} not recognized for E3NN") - else: - raise ValueError(f"Model type {model_type} not recognized") - - _SCORE_MANIFOLD_ALIAS = { ScoreType.chemgauss4: DockingResultCols.DOCKING_SCORE_POSIT.value, ScoreType.FINT: DockingResultCols.FITNESS_SCORE_FINT.value, @@ -452,447 +410,6 @@ def _dispatch(self, inputs: list[Path], **kwargs) -> list[Score]: return self._dispatch(complexes) -class FINTScorer(ScorerBase): - """ - Score using Fitness Interaction Score - - Overloaded to accept DockingResults, Complexes, or Paths to PDB files. - """ - - score_type: ScoreType = Field(ScoreType.FINT, description="Type of score") - units: ClassVar[ScoreUnits.arbitrary] = ScoreUnits.arbitrary - target: TargetTags = Field(..., description="Which target to use for scoring") - - @validator("target") - @classmethod - def validate_target(cls, v): - if not target_has_fitness_data(v): - raise ValueError( - "target does not have fitness data so cannot use FINTScorer" - ) - return v - - @dask_vmap(["inputs"]) - @backend_wrapper("inputs") - def _score( - self, - inputs: Union[list[DockingResult], list[Complex], list[Path]], - return_for_disk_backend: bool = False, - **kwargs, - ) -> list[Score]: - """ - Score the inputs, dispatching based on type. - """ - return self._dispatch( - inputs, return_for_disk_backend=return_for_disk_backend, **kwargs - ) - - @multimethod - def _dispatch( - self, - inputs: list[DockingResult], - return_for_disk_backend: bool = False, - **kwargs, - ) -> list[Score]: - """ - Dispatch for DockingResults - """ - results = [] - for inp in inputs: - _, fint_score = compute_fint_score( - inp.to_protein(), inp.posed_ligand.to_oemol(), self.target - ) - - sc = Score.from_score_and_docking_result( - fint_score, self.score_type, self.units, inp - ) - # overwrite the input with the path to the file - if return_for_disk_backend: - sc.input = _get_disk_path_from_docking_result(inp) - - results.append(sc) - - return results - - @_dispatch.register - def _dispatch(self, inputs: list[Complex], **kwargs): - """ - Dispatch for Complexes - """ - results = [] - for inp in inputs: - _, fint_score = compute_fint_score( - inp.target.to_oemol(), inp.ligand.to_oemol(), self.target - ) - results.append( - Score.from_score_and_complex( - fint_score, self.score_type, self.units, inp - ) - ) - return results - - @_dispatch.register - def _dispatch(self, inputs: list[Path], **kwargs): - """ - Dispatch for PDB files from disk - """ - # assuming reading PDB files from disk - complexes = [ - Complex.from_pdb( - p, - ligand_kwargs={"compound_name": f"{p.stem}_ligand"}, - target_kwargs={"target_name": f"{p.stem}_target"}, - ) - for p in inputs - ] - - return self._dispatch(complexes, **kwargs) - - -# keep track of all the ml scorers -_ml_scorer_classes_meta = [] - - -# decorator to register all the ml scorers -def register_ml_scorer(cls): - _ml_scorer_classes_meta.append(cls) - return cls - - -class MLModelScorer(ScorerBase): - """ - Baseclass to score from some kind of ML model, including 2D or 3D models - """ - - model_type: ClassVar[ModelType.INVALID] = ModelType.INVALID - score_type: ScoreType = Field(..., description="Type of score") - endpoint: Optional[str] = Field(None, description="Endpoint biological property") - units: ClassVar[ScoreUnits.INVALID] = ScoreUnits.INVALID - - targets: Any = Field( - ..., - description="Which targets can this model do predictions for", # FIXME: Optional[set[TargetTags]] - ) - model_name: str = Field(..., description="String indicating which model to use") - inference_cls: InferenceBase = Field(..., description="Inference class") - - @classmethod - def from_latest_by_target(cls, target: TargetTags): - if cls.model_type == ModelType.INVALID: - raise Exception("trying to instantiate some kind a baseclass") - inference_cls = get_inference_cls_from_model_type(cls.model_type) - inference_instance = inference_cls.from_latest_by_target(target) - if inference_instance is None: - logger.warn( - f"no ML model of type {cls.model_type} found for target: {target}, skipping" - ) - return None - else: - try: - instance = cls( - targets=inference_instance.targets, - model_name=inference_instance.model_name, - inference_cls=inference_instance, - endpoint=inference_instance.model_spec.endpoint, - score_type=endpoint_and_model_type_to_score_type( - inference_instance.model_spec.endpoint, cls.model_type - ), - ) - return instance - except Exception as e: - logger.error(f"error instantiating MLModelScorer: {e}") - return None - - @staticmethod - def from_latest_by_target_and_type(target: TargetTags, type: ModelType): - """ - Get the latest ML Scorer by target and type. - - Parameters - ---------- - target : TargetTags - Target to get the scorer for - type : ModelType - Type of model to get the scorer for - """ - if type == ModelType.INVALID: - raise Exception("trying to instantiate some kind a baseclass") - scorer_class = get_ml_scorer_cls_from_model_type(type) - return scorer_class.from_latest_by_target(target) - - @classmethod - def from_model_name(cls, model_name: str): - if cls.model_type == ModelType.INVALID: - raise Exception("trying to instantiate some kind a baseclass") - inference_cls = get_inference_cls_from_model_type(cls.model_type) - inference_instance = inference_cls.from_model_name(model_name) - if inference_instance is None: - logger.warn( - f"no ML model of type {cls.model_type} found for model_name: {model_name}, skipping" - ) - return None - else: - try: - instance = cls( - targets=inference_instance.targets, - model_name=inference_instance.model_name, - inference_cls=inference_instance, - endpoint=inference_instance.model_spec.endpoint, - score_type=endpoint_and_model_type_to_score_type( - inference_instance.model_spec.endpoint, cls.model_type - ), - ) - return instance - except Exception as e: - logger.error(f"error instantiating MLModelScorer: {e}") - return None - - @staticmethod - def load_model_specs( - models: list[MLModelSpecBase], - ) -> list["MLModelScorer"]: # noqa: F821 - """ - Load a list of models into scorers. - - Parameters - ---------- - models : list[MLModelSpecBase] - List of models to load - """ - scorers = [] - for model in models: - scorer_class = get_ml_scorer_cls_from_model_type(model.type) - scorer = scorer_class.from_model_name(model.name) - if scorer is not None: - scorers.append(scorer) - return scorers - - -@register_ml_scorer -class GATScorer(MLModelScorer): - """ - Scoring using GAT ML Model - """ - - model_type: ClassVar[ModelType.GAT] = ModelType.GAT - units: ClassVar[ScoreUnits.pIC50] = ScoreUnits.pIC50 - - @dask_vmap(["inputs"]) - @backend_wrapper("inputs") - def _score( - self, - inputs: Union[list[DockingResult], list[str], list[Ligand]], - return_for_disk_backend: bool = False, - **kwargs, - ) -> list[Score]: - """ - Score the inputs, dispatching based on type. - """ - return self._dispatch( - inputs, return_for_disk_backend=return_for_disk_backend, **kwargs - ) - - @multimethod - def _dispatch( - self, - inputs: list[DockingResult], - return_for_disk_backend: bool = False, - **kwargs, - ) -> list[Score]: - """ - Dispatch for DockingResults - """ - results = [] - for inp in inputs: - gat_score = self.inference_cls.predict_from_smiles(inp.posed_ligand.smiles) - sc = Score.from_score_and_docking_result( - gat_score, - self.score_type, - self.units, - inp, - ) - # overwrite the input with the path to the file - if return_for_disk_backend: - sc.input = _get_disk_path_from_docking_result(inp) - results.append(sc) - return results - - @_dispatch.register - def _dispatch(self, inputs: list[str], **kwargs) -> list[Score]: - """ - Dispatch for SMILES strings - """ - results = [] - for inp in inputs: - gat_score = self.inference_cls.predict_from_smiles(inp) - results.append( - Score.from_score_and_smiles( - gat_score, - inp, - self.score_type, - self.units, - ) - ) - return results - - @_dispatch.register - def _dispatch(self, inputs: list[Ligand], **kwargs) -> list[Score]: - """ - Dispatch for Ligands - """ - results = [] - for inp in inputs: - gat_score = self.inference_cls.predict_from_smiles(inp.smiles) - results.append( - Score.from_score_and_ligand( - gat_score, - inp, - self.score_type, - self.units, - ) - ) - return results - - -class E3MLModelScorer(MLModelScorer): - """ - Scoring using ML Models that operate over 3D structures - These all share an interface so we can use multimethods to dispatch - for the different input types for all subclasses. - """ - - model_type: ClassVar[ModelType.INVALID] = ModelType.INVALID - units: ClassVar[ScoreUnits.INVALID] = ScoreUnits.INVALID - - @dask_vmap(["inputs"]) - @backend_wrapper("inputs") - def _score( - self, - inputs: Union[list[DockingResult], list[Complex], list[Path]], - return_for_disk_backend: bool = False, - **kwargs, - ) -> list[Score]: - return self._dispatch( - inputs, return_for_disk_backend=return_for_disk_backend, **kwargs - ) - - @multimethod - def _dispatch( - self, - inputs: list[DockingResult], - return_for_disk_backend: bool = False, - **kwargs, - ) -> list[Score]: - results = [] - for inp in inputs: - score = self.inference_cls.predict_from_oemol(inp.to_posed_oemol()) - - sc = Score.from_score_and_docking_result( - score, self.score_type, self.units, inp - ) - # overwrite the input with the path to the file - if return_for_disk_backend: - sc.input = _get_disk_path_from_docking_result(inp) - results.append(sc) - - return results - - @_dispatch.register - def _dispatch(self, inputs: list[Complex], **kwargs) -> list[Score]: - results = [] - for inp in inputs: - score = self.inference_cls.predict_from_oemol(inp.to_combined_oemol()) - results.append( - Score.from_score_and_complex(score, self.score_type, self.units, inp) - ) - return results - - @_dispatch.register - def _dispatch(self, inputs: list[Path], **kwargs) -> list[Score]: - # assuming reading PDB files from disk - complexes = [ - Complex.from_pdb( - p, - ligand_kwargs={"compound_name": f"{p.stem}_ligand"}, - target_kwargs={"target_name": f"{p.stem}_target"}, - ) - for i, p in enumerate(inputs) - ] - return self._dispatch(complexes, **kwargs) - - -@register_ml_scorer -class SchnetScorer(E3MLModelScorer): - """ - Scoring using Schnet ML Model - """ - - model_type: ClassVar[ModelType.schnet] = ModelType.schnet - units: ClassVar[ScoreUnits.pIC50] = ScoreUnits.pIC50 - - -@register_ml_scorer -class E3NNScorer(E3MLModelScorer): - """ - Scoring using e3nn ML Model - """ - - model_type: ClassVar[ModelType.e3nn] = ModelType.e3nn - units: ClassVar[ScoreUnits.pIC50] = ScoreUnits.pIC50 - - -def get_ml_scorer_cls_from_model_type(model_type: ModelType): - instantiable_classes = [ - m for m in _ml_scorer_classes_meta if m.model_type != ModelType.INVALID - ] - scorer_class = [m for m in instantiable_classes if m.model_type == model_type] - if len(scorer_class) != 1: - raise Exception("Somehow got multiple scorers") - return scorer_class[0] - - -class MetaScorer(BaseModel): - """ - Score from a combination of other scorers, the scorers must share an input type, - """ - - scorers: list[ScorerBase] = Field(..., description="Scorers to score with") - - def score( - self, - inputs: list[DockingResult], - use_dask: bool = False, - dask_client=None, - failure_mode=FailureMode.SKIP, - backend=BackendType.IN_MEMORY, - reconstruct_cls=None, - return_df: bool = False, - return_for_disk_backend: bool = False, - ) -> list[Score]: - """ - Score the inputs using all the scorers provided in the constructor - """ - results = [] - for scorer in self.scorers: - vals = scorer.score( - inputs=inputs, - use_dask=use_dask, - dask_client=dask_client, - failure_mode=failure_mode, - backend=backend, - reconstruct_cls=reconstruct_cls, - return_df=return_df, - pivot=False, - return_for_disk_backend=return_for_disk_backend, - ) - results.append(vals) - - if return_df: - return Score._combine_and_pivot_scores_df(results) - - return np.ravel(results).tolist() - - class SymClashScorer(ScorerBase): """ Scoring, checking for clashes between ligand and target diff --git a/asapdiscovery-docking/asapdiscovery/docking/tests/test_scorers.py b/asapdiscovery-docking/asapdiscovery/docking/tests/test_scorers.py index 3de820fa..08fa1cd4 100644 --- a/asapdiscovery-docking/asapdiscovery/docking/tests/test_scorers.py +++ b/asapdiscovery-docking/asapdiscovery/docking/tests/test_scorers.py @@ -1,11 +1,6 @@ import pytest from asapdiscovery.docking.scorer import ( ChemGauss4Scorer, - E3NNScorer, - FINTScorer, - GATScorer, - MetaScorer, - SchnetScorer, ) @@ -20,77 +15,3 @@ def test_chemgauss_scorer(use_dask, return_df, data_fixture, request): scorer = ChemGauss4Scorer() scores = scorer.score([data], use_dask=use_dask, return_df=return_df) assert len(scores) == 1 - - -@pytest.mark.parametrize("data_fixture", ["results_simple_nolist", "ligand", "smiles"]) -@pytest.mark.parametrize("return_df", [True, False]) -@pytest.mark.parametrize("use_dask", [True, False]) -def test_gat_scorer(use_dask, return_df, data_fixture, request): - data = request.getfixturevalue(data_fixture) - scorer = GATScorer.from_latest_by_target("SARS-CoV-2-Mpro") - scores = scorer.score([data], use_dask=use_dask, return_df=return_df) - assert len(scores) == 1 - - -@pytest.mark.parametrize( - "data_fixture", ["results_simple_nolist", "complex_simple", "pdb_simple"] -) -@pytest.mark.parametrize("return_df", [True, False]) -@pytest.mark.parametrize("use_dask", [True, False]) -def test_schnet_scorer(use_dask, return_df, data_fixture, request): - data = request.getfixturevalue(data_fixture) - scorer = SchnetScorer.from_latest_by_target("SARS-CoV-2-Mpro") - scores = scorer.score([data], use_dask=use_dask, return_df=return_df) - assert len(scores) == 1 - - -@pytest.mark.parametrize( - "data_fixture", ["results_simple_nolist", "complex_simple", "pdb_simple"] -) -@pytest.mark.parametrize("return_df", [True, False]) -@pytest.mark.parametrize("use_dask", [True, False]) -def test_e3nn_scorer(use_dask, return_df, data_fixture, request): - data = request.getfixturevalue(data_fixture) - scorer = E3NNScorer.from_latest_by_target("SARS-CoV-2-Mpro") - scores = scorer.score([data], use_dask=use_dask, return_df=return_df) - assert len(scores) == 1 - - -@pytest.mark.parametrize("use_dask", [True, False]) -def test_meta_scorer(results, use_dask): - scorer = MetaScorer( - scorers=[ - ChemGauss4Scorer(), - GATScorer.from_latest_by_target("SARS-CoV-2-Mpro"), - SchnetScorer.from_latest_by_target("SARS-CoV-2-Mpro"), - ] - ) - - scores = scorer.score(results, use_dask=use_dask) - assert len(scores) == 3 - - -def test_meta_scorer_df(results_multi): - scorer = MetaScorer( - scorers=[ - ChemGauss4Scorer(), - GATScorer.from_latest_by_target("SARS-CoV-2-Mpro"), - SchnetScorer.from_latest_by_target("SARS-CoV-2-Mpro"), - E3NNScorer.from_latest_by_target("SARS-CoV-2-Mpro"), - ] - ) - - scores = scorer.score(results_multi, return_df=True) - assert len(scores) == 2 # 3 scorers for each of 2 inputs - - -@pytest.mark.parametrize( - "data_fixture", ["results_simple_nolist", "complex_simple", "pdb_simple"] -) -@pytest.mark.parametrize("return_df", [True, False]) -@pytest.mark.parametrize("use_dask", [True, False]) -def test_FINT_scorer(use_dask, return_df, data_fixture, request): - data = request.getfixturevalue(data_fixture) - scorer = FINTScorer(target="SARS-CoV-2-Mpro") - scores = scorer.score([data], use_dask=use_dask, return_df=return_df) - assert len(scores) == 1 diff --git a/asapdiscovery-ml/asapdiscovery/ml/scorers.py b/asapdiscovery-ml/asapdiscovery/ml/scorers.py new file mode 100644 index 00000000..7ffc5152 --- /dev/null +++ b/asapdiscovery-ml/asapdiscovery/ml/scorers.py @@ -0,0 +1,404 @@ +from pathlib import Path +from typing import ClassVar, Optional, Any, Union + +import numpy as np +from mtenn.config import ModelType +from multimethod import multimethod +from pydantic.v1 import Field, BaseModel + +from asapdiscovery.docking.docking import DockingResult +from asapdiscovery.docking.scorer import ( + ScorerBase, + ScoreType, + ScoreUnits, + logger, + Score, + _get_disk_path_from_docking_result, +) +from asapdiscovery.ml.inference import InferenceBase, get_inference_cls_from_model_type +from asapdiscovery.ml.models import MLModelSpecBase +from asapdiscovery.data.schema.complex import Complex +from asapdiscovery.data.schema.ligand import Ligand +from asapdiscovery.data.services.postera.manifold_data_validation import TargetTags +from asapdiscovery.data.util.dask_utils import ( + dask_vmap, + backend_wrapper, + FailureMode, + BackendType, +) + +_ml_scorer_classes_meta = [] + + +def register_ml_scorer(cls): + _ml_scorer_classes_meta.append(cls) + return cls + + +class MLModelScorer(ScorerBase): + """ + Baseclass to score from some kind of ML model, including 2D or 3D models + """ + + model_type: ClassVar[ModelType.INVALID] = ModelType.INVALID + score_type: ScoreType = Field(..., description="Type of score") + endpoint: Optional[str] = Field(None, description="Endpoint biological property") + units: ClassVar[ScoreUnits.INVALID] = ScoreUnits.INVALID + + targets: Any = Field( + ..., + description="Which targets can this model do predictions for", # FIXME: Optional[set[TargetTags]] + ) + model_name: str = Field(..., description="String indicating which model to use") + inference_cls: InferenceBase = Field(..., description="Inference class") + + @classmethod + def from_latest_by_target(cls, target: TargetTags): + if cls.model_type == ModelType.INVALID: + raise Exception("trying to instantiate some kind a baseclass") + inference_cls = get_inference_cls_from_model_type(cls.model_type) + inference_instance = inference_cls.from_latest_by_target(target) + if inference_instance is None: + logger.warn( + f"no ML model of type {cls.model_type} found for target: {target}, skipping" + ) + return None + else: + try: + instance = cls( + targets=inference_instance.targets, + model_name=inference_instance.model_name, + inference_cls=inference_instance, + endpoint=inference_instance.model_spec.endpoint, + score_type=endpoint_and_model_type_to_score_type( + inference_instance.model_spec.endpoint, cls.model_type + ), + ) + return instance + except Exception as e: + logger.error(f"error instantiating MLModelScorer: {e}") + return None + + @staticmethod + def from_latest_by_target_and_type(target: TargetTags, type: ModelType): + """ + Get the latest ML Scorer by target and type. + + Parameters + ---------- + target : TargetTags + Target to get the scorer for + type : ModelType + Type of model to get the scorer for + """ + if type == ModelType.INVALID: + raise Exception("trying to instantiate some kind a baseclass") + scorer_class = get_ml_scorer_cls_from_model_type(type) + return scorer_class.from_latest_by_target(target) + + @classmethod + def from_model_name(cls, model_name: str): + if cls.model_type == ModelType.INVALID: + raise Exception("trying to instantiate some kind a baseclass") + inference_cls = get_inference_cls_from_model_type(cls.model_type) + inference_instance = inference_cls.from_model_name(model_name) + if inference_instance is None: + logger.warn( + f"no ML model of type {cls.model_type} found for model_name: {model_name}, skipping" + ) + return None + else: + try: + instance = cls( + targets=inference_instance.targets, + model_name=inference_instance.model_name, + inference_cls=inference_instance, + endpoint=inference_instance.model_spec.endpoint, + score_type=endpoint_and_model_type_to_score_type( + inference_instance.model_spec.endpoint, cls.model_type + ), + ) + return instance + except Exception as e: + logger.error(f"error instantiating MLModelScorer: {e}") + return None + + @staticmethod + def load_model_specs( + models: list[MLModelSpecBase], + ) -> list["MLModelScorer"]: # noqa: F821 + """ + Load a list of models into scorers. + + Parameters + ---------- + models : list[MLModelSpecBase] + List of models to load + """ + scorers = [] + for model in models: + scorer_class = get_ml_scorer_cls_from_model_type(model.type) + scorer = scorer_class.from_model_name(model.name) + if scorer is not None: + scorers.append(scorer) + return scorers + + +@register_ml_scorer +class GATScorer(MLModelScorer): + """ + Scoring using GAT ML Model + """ + + model_type: ClassVar[ModelType.GAT] = ModelType.GAT + units: ClassVar[ScoreUnits.pIC50] = ScoreUnits.pIC50 + + @dask_vmap(["inputs"]) + @backend_wrapper("inputs") + def _score( + self, + inputs: Union[list[DockingResult], list[str], list[Ligand]], + return_for_disk_backend: bool = False, + **kwargs, + ) -> list[Score]: + """ + Score the inputs, dispatching based on type. + """ + return self._dispatch( + inputs, return_for_disk_backend=return_for_disk_backend, **kwargs + ) + + @multimethod + def _dispatch( + self, + inputs: list[DockingResult], + return_for_disk_backend: bool = False, + **kwargs, + ) -> list[Score]: + """ + Dispatch for DockingResults + """ + results = [] + for inp in inputs: + gat_score = self.inference_cls.predict_from_smiles(inp.posed_ligand.smiles) + sc = Score.from_score_and_docking_result( + gat_score, + self.score_type, + self.units, + inp, + ) + # overwrite the input with the path to the file + if return_for_disk_backend: + sc.input = _get_disk_path_from_docking_result(inp) + results.append(sc) + return results + + @_dispatch.register + def _dispatch(self, inputs: list[str], **kwargs) -> list[Score]: + """ + Dispatch for SMILES strings + """ + results = [] + for inp in inputs: + gat_score = self.inference_cls.predict_from_smiles(inp) + results.append( + Score.from_score_and_smiles( + gat_score, + inp, + self.score_type, + self.units, + ) + ) + return results + + @_dispatch.register + def _dispatch(self, inputs: list[Ligand], **kwargs) -> list[Score]: + """ + Dispatch for Ligands + """ + results = [] + for inp in inputs: + gat_score = self.inference_cls.predict_from_smiles(inp.smiles) + results.append( + Score.from_score_and_ligand( + gat_score, + inp, + self.score_type, + self.units, + ) + ) + return results + + +class E3MLModelScorer(MLModelScorer): + """ + Scoring using ML Models that operate over 3D structures + These all share an interface so we can use multimethods to dispatch + for the different input types for all subclasses. + """ + + model_type: ClassVar[ModelType.INVALID] = ModelType.INVALID + units: ClassVar[ScoreUnits.INVALID] = ScoreUnits.INVALID + + @dask_vmap(["inputs"]) + @backend_wrapper("inputs") + def _score( + self, + inputs: Union[list[DockingResult], list[Complex], list[Path]], + return_for_disk_backend: bool = False, + **kwargs, + ) -> list[Score]: + return self._dispatch( + inputs, return_for_disk_backend=return_for_disk_backend, **kwargs + ) + + @multimethod + def _dispatch( + self, + inputs: list[DockingResult], + return_for_disk_backend: bool = False, + **kwargs, + ) -> list[Score]: + results = [] + for inp in inputs: + score = self.inference_cls.predict_from_oemol(inp.to_posed_oemol()) + + sc = Score.from_score_and_docking_result( + score, self.score_type, self.units, inp + ) + # overwrite the input with the path to the file + if return_for_disk_backend: + sc.input = _get_disk_path_from_docking_result(inp) + results.append(sc) + + return results + + @_dispatch.register + def _dispatch(self, inputs: list[Complex], **kwargs) -> list[Score]: + results = [] + for inp in inputs: + score = self.inference_cls.predict_from_oemol(inp.to_combined_oemol()) + results.append( + Score.from_score_and_complex(score, self.score_type, self.units, inp) + ) + return results + + @_dispatch.register + def _dispatch(self, inputs: list[Path], **kwargs) -> list[Score]: + # assuming reading PDB files from disk + complexes = [ + Complex.from_pdb( + p, + ligand_kwargs={"compound_name": f"{p.stem}_ligand"}, + target_kwargs={"target_name": f"{p.stem}_target"}, + ) + for i, p in enumerate(inputs) + ] + return self._dispatch(complexes, **kwargs) + + +@register_ml_scorer +class SchnetScorer(E3MLModelScorer): + """ + Scoring using Schnet ML Model + """ + + model_type: ClassVar[ModelType.schnet] = ModelType.schnet + units: ClassVar[ScoreUnits.pIC50] = ScoreUnits.pIC50 + + +@register_ml_scorer +class E3NNScorer(E3MLModelScorer): + """ + Scoring using e3nn ML Model + """ + + model_type: ClassVar[ModelType.e3nn] = ModelType.e3nn + units: ClassVar[ScoreUnits.pIC50] = ScoreUnits.pIC50 + + +def get_ml_scorer_cls_from_model_type(model_type: ModelType): + instantiable_classes = [ + m for m in _ml_scorer_classes_meta if m.model_type != ModelType.INVALID + ] + scorer_class = [m for m in instantiable_classes if m.model_type == model_type] + if len(scorer_class) != 1: + raise Exception("Somehow got multiple scorers") + return scorer_class[0] + + +class MetaScorer(BaseModel): + """ + Score from a combination of other scorers, the scorers must share an input type, + """ + + scorers: list[ScorerBase] = Field(..., description="Scorers to score with") + + def score( + self, + inputs: list[DockingResult], + use_dask: bool = False, + dask_client=None, + failure_mode=FailureMode.SKIP, + backend=BackendType.IN_MEMORY, + reconstruct_cls=None, + return_df: bool = False, + return_for_disk_backend: bool = False, + ) -> list[Score]: + """ + Score the inputs using all the scorers provided in the constructor + """ + results = [] + for scorer in self.scorers: + vals = scorer.score( + inputs=inputs, + use_dask=use_dask, + dask_client=dask_client, + failure_mode=failure_mode, + backend=backend, + reconstruct_cls=reconstruct_cls, + return_df=return_df, + pivot=False, + return_for_disk_backend=return_for_disk_backend, + ) + results.append(vals) + + if return_df: + return Score._combine_and_pivot_scores_df(results) + + return np.ravel(results).tolist() + + +def endpoint_and_model_type_to_score_type(endpoint: str, model_type: str) -> ScoreType: + """ + Convert an endpoint to a score type. + + Parameters + ---------- + endpoint : str + Endpoint to convert + + Returns + ------- + ScoreType + Score type + """ + if model_type == ModelType.GAT: + if endpoint == "pIC50": # TODO: make this an enum + return ScoreType.GAT_pIC50 + elif endpoint == "LogD": + return ScoreType.GAT_LogD + else: + raise ValueError(f"Endpoint {endpoint} not recognized, for GAT") + elif model_type == ModelType.schnet: + if endpoint == "pIC50": + return ScoreType.schnet_pIC50 + else: + raise ValueError(f"Endpoint {endpoint} not recognized, for Schnet") + elif model_type == ModelType.e3nn: + if endpoint == "pIC50": + return ScoreType.e3nn_pIC50 + else: + raise ValueError(f"Endpoint {endpoint} not recognized for E3NN") + else: + raise ValueError(f"Model type {model_type} not recognized") diff --git a/asapdiscovery-ml/asapdiscovery/ml/tests/test_scorers.py b/asapdiscovery-ml/asapdiscovery/ml/tests/test_scorers.py new file mode 100644 index 00000000..eb28af92 --- /dev/null +++ b/asapdiscovery-ml/asapdiscovery/ml/tests/test_scorers.py @@ -0,0 +1,66 @@ +import pytest + +from asapdiscovery.docking.scorer import ChemGauss4Scorer +from asapdiscovery.ml.scorers import GATScorer, SchnetScorer, E3NNScorer, MetaScorer + + +@pytest.mark.parametrize("data_fixture", ["results_simple_nolist", "ligand", "smiles"]) +@pytest.mark.parametrize("return_df", [True, False]) +@pytest.mark.parametrize("use_dask", [True, False]) +def test_gat_scorer(use_dask, return_df, data_fixture, request): + data = request.getfixturevalue(data_fixture) + scorer = GATScorer.from_latest_by_target("SARS-CoV-2-Mpro") + scores = scorer.score([data], use_dask=use_dask, return_df=return_df) + assert len(scores) == 1 + + +@pytest.mark.parametrize( + "data_fixture", ["results_simple_nolist", "complex_simple", "pdb_simple"] +) +@pytest.mark.parametrize("return_df", [True, False]) +@pytest.mark.parametrize("use_dask", [True, False]) +def test_schnet_scorer(use_dask, return_df, data_fixture, request): + data = request.getfixturevalue(data_fixture) + scorer = SchnetScorer.from_latest_by_target("SARS-CoV-2-Mpro") + scores = scorer.score([data], use_dask=use_dask, return_df=return_df) + assert len(scores) == 1 + + +@pytest.mark.parametrize( + "data_fixture", ["results_simple_nolist", "complex_simple", "pdb_simple"] +) +@pytest.mark.parametrize("return_df", [True, False]) +@pytest.mark.parametrize("use_dask", [True, False]) +def test_e3nn_scorer(use_dask, return_df, data_fixture, request): + data = request.getfixturevalue(data_fixture) + scorer = E3NNScorer.from_latest_by_target("SARS-CoV-2-Mpro") + scores = scorer.score([data], use_dask=use_dask, return_df=return_df) + assert len(scores) == 1 + + +@pytest.mark.parametrize("use_dask", [True, False]) +def test_meta_scorer(results, use_dask): + scorer = MetaScorer( + scorers=[ + ChemGauss4Scorer(), + GATScorer.from_latest_by_target("SARS-CoV-2-Mpro"), + SchnetScorer.from_latest_by_target("SARS-CoV-2-Mpro"), + ] + ) + + scores = scorer.score(results, use_dask=use_dask) + assert len(scores) == 3 + + +def test_meta_scorer_df(results_multi): + scorer = MetaScorer( + scorers=[ + ChemGauss4Scorer(), + GATScorer.from_latest_by_target("SARS-CoV-2-Mpro"), + SchnetScorer.from_latest_by_target("SARS-CoV-2-Mpro"), + E3NNScorer.from_latest_by_target("SARS-CoV-2-Mpro"), + ] + ) + + scores = scorer.score(results_multi, return_df=True) + assert len(scores) == 2 # 3 scorers for each of 2 inputs diff --git a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/cross_docking.py b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/cross_docking.py index 01340fdc..19171a24 100644 --- a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/cross_docking.py +++ b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/cross_docking.py @@ -20,7 +20,8 @@ ) from asapdiscovery.docking.docking_data_validation import DockingResultCols from asapdiscovery.docking.openeye import POSIT_METHOD, POSIT_RELAX_MODE, POSITDocker -from asapdiscovery.docking.scorer import ChemGauss4Scorer, MetaScorer +from asapdiscovery.docking.scorer import ChemGauss4Scorer +from asapdiscovery.ml.scorers import MetaScorer from asapdiscovery.modeling.protein_prep import ProteinPrepper from asapdiscovery.workflows.docking_workflows.workflows import ( DockingWorkflowInputsBase, diff --git a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/large_scale_docking.py b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/large_scale_docking.py index 1d91053f..6abbcdd9 100644 --- a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/large_scale_docking.py +++ b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/large_scale_docking.py @@ -33,10 +33,9 @@ from asapdiscovery.docking.openeye import POSITDocker from asapdiscovery.docking.scorer import ( ChemGauss4Scorer, - FINTScorer, - MetaScorer, - MLModelScorer, ) +from asapdiscovery.dataviz.fint_score import FINTScorer +from asapdiscovery.ml.scorers import MLModelScorer, MetaScorer from asapdiscovery.ml.models import ASAPMLModelRegistry from asapdiscovery.modeling.protein_prep import ProteinPrepper from asapdiscovery.spectrum.fitness import target_has_fitness_data diff --git a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/ligand_transfer_docking.py b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/ligand_transfer_docking.py index 73d83c0b..585f10c1 100644 --- a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/ligand_transfer_docking.py +++ b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/ligand_transfer_docking.py @@ -28,7 +28,8 @@ from asapdiscovery.docking.docking import write_results_to_multi_sdf from asapdiscovery.docking.docking_data_validation import DockingResultCols from asapdiscovery.docking.openeye import POSIT_METHOD, POSIT_RELAX_MODE, POSITDocker -from asapdiscovery.docking.scorer import ChemGauss4Scorer, MetaScorer, MLModelScorer +from asapdiscovery.docking.scorer import ChemGauss4Scorer +from asapdiscovery.ml.scorers import MLModelScorer, MetaScorer from asapdiscovery.ml.models import ASAPMLModelRegistry from asapdiscovery.modeling.protein_prep import LigandTransferProteinPrepper from asapdiscovery.simulation.simulate import OpenMMPlatform, VanillaMDSimulator diff --git a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/small_scale_docking.py b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/small_scale_docking.py index d9d080b6..3e5ae5f1 100644 --- a/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/small_scale_docking.py +++ b/asapdiscovery-workflows/asapdiscovery/workflows/docking_workflows/small_scale_docking.py @@ -39,10 +39,9 @@ from asapdiscovery.docking.openeye import POSITDocker from asapdiscovery.docking.scorer import ( ChemGauss4Scorer, - FINTScorer, - MetaScorer, - MLModelScorer, ) +from asapdiscovery.dataviz.fint_score import FINTScorer +from asapdiscovery.ml.scorers import MLModelScorer, MetaScorer from asapdiscovery.ml.models import ASAPMLModelRegistry from asapdiscovery.modeling.protein_prep import ProteinPrepper from asapdiscovery.simulation.simulate import OpenMMPlatform, VanillaMDSimulator