Skip to content

Commit

Permalink
mol2any: add mol2bool element (#59)
Browse files Browse the repository at this point in the history
- Add new MolToBool element that converts values to
      a bool array. Can be used for masking.
  • Loading branch information
JochenSiegWork authored Aug 15, 2024
1 parent 253eaea commit 634ecc4
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 0 deletions.
52 changes: 52 additions & 0 deletions molpipeline/mol2any/mol2bool.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
"""Pipeline elements for converting instances to bool."""

from typing import Any

from molpipeline.abstract_pipeline_elements.core import (
MolToAnyPipelineElement,
InvalidInstance,
)


class MolToBool(MolToAnyPipelineElement):
"""Element to generate a bool array from input."""

def __init__(
self,
name: str = "Mol2Bool",
n_jobs: int = 1,
uuid: str | None = None,
) -> None:
"""Initialize MolToBinaryPipelineElement.
Parameters
----------
name: str, optional (default="Mol2Bool")
name of PipelineElement
n_jobs: int, optional (default=1)
number of jobs to use for parallelization
uuid: Optional[str], optional (default=None)
uuid of PipelineElement, by default None
Returns
-------
None
"""
super().__init__(name=name, n_jobs=n_jobs, uuid=uuid)

def pretransform_single(self, value: Any) -> bool:
"""Transform a value to a bool representation.
Parameters
----------
value: Any
Value to be transformed to bool representation.
Returns
-------
str
Binary representation of molecule.
"""
if isinstance(value, InvalidInstance):
return False
return True
24 changes: 24 additions & 0 deletions tests/test_elements/test_mol2any/test_mol2bool.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
"""Test mol to bool conversion."""

import unittest

from molpipeline.abstract_pipeline_elements.core import InvalidInstance
from molpipeline.mol2any.mol2bool import MolToBool


class TestMolToBool(unittest.TestCase):
"""Unittest for MolToBool."""

def test_bool_conversion(self) -> None:
"""Test if the invalid instances are converted to bool."""

mol2bool = MolToBool()
result = mol2bool.transform(
[
1,
2,
InvalidInstance(element_id="test", message="test", element_name="Test"),
4,
]
)
self.assertEqual(result, [True, True, False, True])

0 comments on commit 634ecc4

Please sign in to comment.