separate docking scorers into separate files

asapdiscovery-docking/asapdiscovery/docking/fint_scorer.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)

asapdiscovery-docking/asapdiscovery/docking/meta_scorer.py

@@ -0,0 +1,48 @@
+import numpy as np
+from pydantic.v1 import BaseModel, Field
+
+from asapdiscovery.docking.docking import DockingResult
+from asapdiscovery.docking.scorer import ScorerBase, Score
+from asapdiscovery.data.util.dask_utils import FailureMode, BackendType
+
+
+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()