Filter update

molpipeline/abstract_pipeline_elements/mol2mol/__init__.py

@@ -0,0 +1,8 @@
+"""Initialize the module for abstract mol2mol elements."""
+
+from molpipeline.abstract_pipeline_elements.mol2mol.filter import (
+    BaseKeepMatchesFilter,
+    BasePatternsFilter,
+)
+
+__all__ = ["BasePatternsFilter", "BaseKeepMatchesFilter"]

molpipeline/abstract_pipeline_elements/mol2mol/filter.py

@@ -0,0 +1,196 @@
+"""Abstract classes for filters."""
+
+import abc
+from typing import Any, Literal, Optional, Union
+
+try:
+    from typing import Self  # type: ignore[attr-defined]
+except ImportError:
+    from typing_extensions import Self
+
+from molpipeline.abstract_pipeline_elements.core import MolToMolPipelineElement
+from molpipeline.utils.value_conversions import count_value_to_tuple
+
+
+class BaseKeepMatchesFilter(MolToMolPipelineElement, abc.ABC):
+    """Filter to keep or remove molecules based on patterns."""
+
+    keep_matches: bool
+    mode: Literal["any", "all"]
+
+    def __init__(
+        self,
+        keep_matches: bool = True,
+        mode: Literal["any", "all"] = "any",
+        name: Optional[str] = None,
+        n_jobs: int = 1,
+        uuid: Optional[str] = None,
+    ) -> None:
+        """Initialize BasePatternsFilter.
+
+        Parameters
+        ----------
+        keep_matches: bool, optional (default: True)
+            If True, molecules containing the specified patterns are kept, else removed.
+        mode: Literal["any", "all"], optional (default: "any")
+            If "any", at least one of the specified patterns must be present in the molecule.
+            If "all", all of the specified patterns must be present in the molecule.
+        name: Optional[str], optional (default: None)
+            Name of the pipeline element.
+        n_jobs: int, optional (default: 1)
+            Number of parallel jobs to use.
+        uuid: str, optional (default: None)
+            Unique identifier of the pipeline element.
+        """
+        super().__init__(name=name, n_jobs=n_jobs, uuid=uuid)
+        self.keep_matches = keep_matches
+        self.mode = mode
+
+    def set_params(self, **parameters: Any) -> Self:
+        """Set parameters of BaseKeepMatchesFilter.
+
+        Parameters
+        ----------
+        parameters: Any
+            Parameters to set.
+
+        Returns
+        -------
+        Self
+            Self.
+        """
+        parameter_copy = dict(parameters)
+        if "keep_matches" in parameter_copy:
+            self.keep_matches = parameter_copy.pop("keep_matches")
+        if "mode" in parameter_copy:
+            self.mode = parameter_copy.pop("mode")
+        super().set_params(**parameter_copy)
+        return self
+
+    def get_params(self, deep: bool = True) -> dict[str, Any]:
+        """Get parameters of PatternFilter.
+
+        Parameters
+        ----------
+        deep: bool, optional (default: True)
+            If True, return the parameters of all subobjects that are PipelineElements.
+
+        Returns
+        -------
+        dict[str, Any]
+            Parameters of BaseKeepMatchesFilter.
+        """
+        params = super().get_params(deep=deep)
+        params["keep_matches"] = self.keep_matches
+        params["mode"] = self.mode
+        return params
+
+
+class BasePatternsFilter(BaseKeepMatchesFilter, abc.ABC):
+    """Filter to keep or remove molecules based on patterns."""
+
+    _patterns: dict[str, tuple[Optional[int], Optional[int]]]
+
+    def __init__(
+        self,
+        patterns: Union[
+            list[str], dict[str, Union[int, tuple[Optional[int], Optional[int]]]]
+        ],
+        keep_matches: bool = True,
+        mode: Literal["any", "all"] = "any",
+        name: Optional[str] = None,
+        n_jobs: int = 1,
+        uuid: Optional[str] = None,
+    ) -> None:
+        """Initialize BasePatternsFilter.
+
+        Parameters
+        ----------
+        patterns: Union[list[str], dict[str, Union[int, tuple[Optional[int], Optional[int]]]]]
+            List of patterns to allow in molecules.
+            Alternatively, a dictionary can be passed with patterns as keys
+            and an int for exact count or a tuple of minimum and maximum.
+        keep_matches: bool, optional (default: True)
+            If True, molecules containing the specified patterns are kept, else removed.
+        mode: Literal["any", "all"], optional (default: "any")
+            If "any", at least one of the specified patterns must be present in the molecule.
+            If "all", all of the specified patterns must be present in the molecule.
+        name: Optional[str], optional (default: None)
+            Name of the pipeline element.
+        n_jobs: int, optional (default: 1)
+            Number of parallel jobs to use.
+        uuid: str, optional (default: None)
+            Unique identifier of the pipeline element.
+        """
+        super().__init__(
+            keep_matches=keep_matches, mode=mode, name=name, n_jobs=n_jobs, uuid=uuid
+        )
+        self.patterns = patterns  # type: ignore
+
+    @property
+    def patterns(self) -> dict[str, tuple[Optional[int], Optional[int]]]:
+        """Get allowed patterns as dict."""
+        return self._patterns
+
+    @patterns.setter
+    def patterns(
+        self,
+        patterns: Union[
+            list[str], dict[str, Union[int, tuple[Optional[int], Optional[int]]]]
+        ],
+    ) -> None:
+        """Set allowed patterns as dict.
+
+        Parameters
+        ----------
+        patterns: Union[list[str], dict[str, Union[int, tuple[Optional[int], Optional[int]]]]]
+            List of patterns.
+        """
+        if isinstance(patterns, (list, set)):
+            self._patterns = {pat: (1, None) for pat in patterns}
+        else:
+            self._patterns = {
+                pat: count_value_to_tuple(count) for pat, count in patterns.items()
+            }
+
+    def get_params(self, deep: bool = True) -> dict[str, Any]:
+        """Get parameters of PatternFilter.
+
+        Parameters
+        ----------
+        deep: bool, optional (default: True)
+            If True, return the parameters of all subobjects that are PipelineElements.
+
+        Returns
+        -------
+        dict[str, Any]
+            Parameters of PatternFilter.
+        """
+        params = super().get_params(deep=deep)
+        if deep:
+            params["patterns"] = {
+                pat: (count_tuple[0], count_tuple[1])
+                for pat, count_tuple in self.patterns.items()
+            }
+        else:
+            params["patterns"] = self.patterns
+        return params
+
+    def set_params(self, **parameters: Any) -> Self:
+        """Set parameters of PatternFilter.
+
+        Parameters
+        ----------
+        parameters: Any
+            Parameters to set.
+
+        Returns
+        -------
+        Self
+            Self.
+        """
+        parameter_copy = dict(parameters)
+        if "patterns" in parameter_copy:
+            self.patterns = parameter_copy.pop("patterns")
+        super().set_params(**parameter_copy)
+        return self

molpipeline/mol2mol/__init__.py

@@ -1,10 +1,13 @@
 """Init the module for mol2mol pipeline elements."""
 
 from molpipeline.mol2mol.filter import (
+    DescriptorsFilter,
     ElementFilter,
     EmptyMoleculeFilter,
     InorganicsFilter,
     MixtureFilter,
+    SmartsFilter,
+    SmilesFilter,
 )
 from molpipeline.mol2mol.reaction import MolToMolReaction
 from molpipeline.mol2mol.scaffolds import MakeScaffoldGeneric, MurckoScaffold
@@ -41,4 +44,7 @@
     "SolventRemover",
     "Uncharger",
     "InorganicsFilter",
+    "SmartsFilter",
+    "SmilesFilter",
+    "DescriptorsFilter",
 )