How to define a new Component

The gufe Component is a GufeTokenizable intended to be used as an extensible point of the library, such that you can describe a simulated system in terms of GufeTokenizables that are compatible with gufe and the rest of the OpenFE ecosystem.

Step 1: Choose a Component to Extend

In many cases, you will likely want your custom Component to inherit from one of the following extensible points that themselves inherit from the Component base class - meaning that all of the following are Components, but come with additional functionality:

In the rare case where you want as much custom implementation as possible, you could inherit directly from Component itself. However, many OpenFE Protocols require specific Component types, in which case sub-classing from a compatible class is recommended.

Step 2: Write Tests for Expected Behavior

As when defining any new GufeTokenizable, you are encouraged to use the GufeTokenizableTestsMixin pytest fixture to ensure that your new Component works as intended.

Note

A key benefit of inheriting from a fully implemented class, such as SmallMoleculeComponent in this example, is that you only need to implement new functionality.

For example:

custom_component.py
from gufe import SmallMoleculeComponent

# this is all the code you need to get the GufeTokenizableTestsMixin to pass,
# since SmallMoleculeComponent fully implemented
class CustomComponent(SmallMoleculeComponent):
    pass
test_custom_component.py
from gufe.tests import GufeTokenizableTestsMixin
from .custom_component import CustomComponent
import pytest
from rdkit import Chem


class TestCustomComponent(GufeTokenizableTestsMixin):
    cls = CustomComponent
    repr = "CustomComponent(name=ethane)"

    @pytest.fixture()
    def instance(self):
        mol = Chem.AddHs(Chem.MolFromSmiles("CC"))
        Chem.AllChem.Compute2DCoords(mol)
        return CustomComponent(rdkit=mol, name="ethane")

Step 3: Define Required Methods

When inheriting from abstract base classes, such as Component, you will need to define anything that is an abstractmethod. This includes both in Component itself, as well as any abstractmethods it inherits from GufeTokenizable (since component is a subclass of GufeTokenizable).

In other cases, such as when inheriting from SmallMoleculeComponent, you will only need to define methods specifically not implemented - in this case _to_dict() and _from_dict().

Also note any base requirements of the parent class from which your Component inherits - for example, anything inheriting from the ExplicitMoleculeComponent base class must be compatible with its to_rdkit and from_rdkit methods.

Step 4: Define Additional Functionality

While the code in Step 2 is technically correct, it doesn’t actually add anything new; it merely creates a new class identical to SmallMoleculeComponent with a new name. To add functionality in addition to SmallMoleculeComponent’s existing functionality, you can add new attributes, such as custom_attribute, and new methods, such as print_custom_attribute:

custom_component.py
from gufe import SmallMoleculeComponent
from rdkit import Chem


class CustomComponent(SmallMoleculeComponent):
    def __init__(self, rdkit: Chem.rdchem.Mol, name: str = "", custom_attribute: int = 4):
        self.custom_attribute = custom_attribute
        super().__init__(rdkit=rdkit, name=name)

    def custom_functionality(self) -> str:
        return f"my custom attribute is {self.custom_attribute}"

    # Since we added a new attribute, must include that attribute in serialization
    # by defining _to_dict and _from_dict
    def _to_dict(self):
        # first, use the parent classes' implementation
        d = super()._to_dict()

        # now, add our custom attribute
        d["custom_attribute"] = self.custom_attribute
        return d

    @classmethod
    def _from_dict(cls, d:dict):
        # first, use the parent classes' implementation
        obj = super()._from_dict(d)
        # now, pass through to construct our custom attribute
        return cls(rdkit=obj._rdkit, name=obj.name, custom_attribute=d.get("custom_attribute"))

Just make sure you test all of your new features!

test_custom_component.py
from gufe.tests import GufeTokenizableTestsMixin
from .custom_component import CustomComponent
import pytest
from rdkit import Chem


class TestCustomComponent(GufeTokenizableTestsMixin):
    cls = CustomComponent
    repr = "CustomComponent(name=ethane)"

    @pytest.fixture()
    def instance(self):
        mol = Chem.AddHs(Chem.MolFromSmiles("CC"))
        Chem.AllChem.Compute2DCoords(mol)

        # it's important we test defining the custom attribute to make sure it round-trips correctly
        test_instance = CustomComponent(rdkit=mol, name="ethane", custom_attribute=8)
        return test_instance

    def test_print_custom_attribute_default(self, instance):
        # test using the instance fixture created above
        assert instance.print_custom_attribute() == "my custom attribute is 4"

    def test_print_custom_attribute_user_defined(self):
        # test by creating a new instance
        mol = Chem.AddHs(Chem.MolFromSmiles("CCC"))
        Chem.AllChem.Compute2DCoords(mol)
        custom_component = CustomComponent(
            rdkit=mol, name="propane", custom_attribute=7
        )
        assert custom_component.print_custom_attribute() == "my custom attribute is 7"