Defining a Custom BioMol Class

In this tutorial, you will

• Learn how to create a custom class that inherits from the BioMol class.

• Define a custom View for feature type-hinting.

• Understand the benefits of using type hints in your code.

The Problem: Lack of Autocompletion

Before we begin, let’s discuss why type-hinting is important. Because the BioMol class doesn’t predefine specific features, it’s hard to know what features are available in a BioMol instance. For example, autocompletion for mol.atoms.positions doesn’t work.

To address this limitation, you will:

1. Define a custom View for feature type-hinting.

2. Create a custom class that inherits from BioMol and uses the defined View.

Let’s start by importing the necessary modules.

import numpy as np

from biomol import BioMol
from biomol.core import FeatureContainer, IndexTable, NodeFeature, View

Preparing container and index table

As in the previous tutorial, prepare a FeatureContainer and an IndexTable to hold features and establish hierarchical relationships.

atom_positions = NodeFeature(
    value=np.array(
        [
            [0.0, 0.0, 0.0],
            [1.4, 0.0, 0.0],
            [1.4, 1.4, 0.0],
            [0.0, 1.4, 0.0],  # ALA-1
            [2.8, 0.0, 0.0],
            [4.2, 0.0, 0.0],
            [4.2, 1.4, 0.0],
            [2.8, 1.4, 0.0],  # GLY-2
            [5.6, 0.0, 0.0],
            [7.0, 0.0, 0.0],
            [7.0, 1.4, 0.0],
            [5.6, 1.4, 0.0],  # ALA-3
        ],
    ),
)
atom_names = NodeFeature(value=np.array(["N", "CA", "C", "O"] * 3))

residue_ids = NodeFeature(value=np.array([1, 2, 3]))
residue_names = NodeFeature(value=np.array(["ALA", "GLY", "ALA"]))

chain_ids = NodeFeature(value=np.array(["A"]))
chain_entities = NodeFeature(value=np.array(["PROTEIN"]))

atom_container = FeatureContainer(
    {"positions": atom_positions, "name": atom_names},
)

residue_container = FeatureContainer(
    {"id": residue_ids, "name": residue_names},
)

chain_container = FeatureContainer(
    {"id": chain_ids, "entity": chain_entities},
)

index_table = IndexTable.from_parents(
    atom_to_res=np.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]),
    res_to_chain=np.array([0, 0, 0]),
)

Defining a Custom BioMol

Next, define a custom View and a custom class that inherits from BioMol class. The View specifies which features are exposed on the FeatureContainer via property methods.

Note

View is used solely for static type-hinting. It does not affect runtime behavior and therefore cannot guarantee that features actually exist at runtime.

class MyAtomView(View["MyAtomView", "MyResidueView", "MyChainView", "MyBioMol"]):
    """Custom atom-level View."""

    @property
    def positions(self) -> NodeFeature:
        """XYZ coordinates of atoms."""

    @property
    def name(self) -> NodeFeature:
        """Atom name."""


class MyResidueView(View["MyAtomView", "MyResidueView", "MyChainView", "MyBioMol"]):
    """Custom residue-level View."""

    @property
    def id(self) -> NodeFeature:
        """Residue ID."""

    @property
    def name(self) -> NodeFeature:
        """Residue name."""


class MyChainView(View["MyAtomView", "MyResidueView", "MyChainView", "MyBioMol"]):
    """Custom chain-level View."""

    @property
    def id(self) -> NodeFeature:
        """Chain ID."""

    @property
    def entity(self) -> NodeFeature:
        """Chain entity type."""


class MyBioMol(BioMol["MyAtomView", "MyResidueView", "MyChainView"]):
    """Custom BioMol class."""

Now, you can create an instance of your custom BioMol class and access its features with type-hinting support.

mol = MyBioMol(
    atom_container=atom_container,
    residue_container=residue_container,
    chain_container=chain_container,
    index_table=index_table,
)

# ruff: noqa: B018

mol.atoms.positions
mol.atoms.name

mol.residues.id
mol.residues.name

mol.chains.id
mol.chains.entity

NodeFeature(value=array(['PROTEIN'], dtype='<U7'))

Type hints remain effective even in complex expressions that involve multiple hierarchical traversals.

mol.atoms.residues.chains.id
mol.atoms.chains.residues.atoms.name
mol.chains.atoms.chains.residues.id

NodeFeature(value=array([1, 2, 3]))