Wrappers for extension¶

In this tutorial, you'll learn about the Wrapper system in MolPy, which provides a composition-based approach to extending functionality.

Introduction¶

Wrappers in MolPy follow a composition over inheritance design pattern. They allow you to extend the functionality of Struct objects (like Atomistic) without modifying the base classes.

Key concepts:

• Wrapper: A container that holds an inner Struct object
• Explicit forwarding: Wrappers explicitly forward selected methods/properties
• Type safety: Uses generics to preserve type information

Let's start by importing MolPy:

In [1]:

import numpy as np

import molpy as mp
from molpy.core.atomistic import Atomistic
from molpy.core.wrappers.base import Wrapper

import numpy as np import molpy as mp from molpy.core.atomistic import Atomistic from molpy.core.wrappers.base import Wrapper

Built-in Wrappers¶

MolPy provides some built-in wrappers. Let's look at the Monomer wrapper:

In [2]:

from molpy.core.wrappers.monomer import Monomer

struct = Atomistic()
struct.add_atoms(
    [
        mp.Atom(x=0.0, y=0.0, z=3.0, element="C"),
        mp.Atom(x=3.0, y=0.0, z=0.0, element="C"),
        mp.Atom(x=0.0, y=3.0, z=0.0, element="C"),
    ]
)

# Wrap it as a monomer
monomer = Monomer(struct, name="ethylene")
print(f"Monomer name: {monomer['name']}")
print(f"Number of atoms: {len(monomer.inner.atoms)}")

from molpy.core.wrappers.monomer import Monomer struct = Atomistic() struct.add_atoms( [ mp.Atom(x=0.0, y=0.0, z=3.0, element="C"), mp.Atom(x=3.0, y=0.0, z=0.0, element="C"), mp.Atom(x=0.0, y=3.0, z=0.0, element="C"), ] ) # Wrap it as a monomer monomer = Monomer(struct, name="ethylene") print(f"Monomer name: {monomer['name']}") print(f"Number of atoms: {len(monomer.inner.atoms)}")

Monomer name: ethylene
Number of atoms: 3

Creating Custom Wrappers¶

You can create your own wrappers by subclassing Wrapper:

In [3]:

class LabeledMolecule(Wrapper[Atomistic]):
    """A custom wrapper that adds a label to a molecule."""

    def __init__(self, inner: Atomistic, label: str):
        super().__init__(inner)
        self.label = label

    @property
    def n_atoms(self) -> int:
        """Explicitly forward the number of atoms."""
        return len(self.inner.atoms)

    @property
    def center_of_mass(self):
        """Calculate center of mass."""
        # This is a simplified example
        atoms = self.inner.atoms
        xyz = np.array([[atom["x"], atom["y"], atom["z"]] for atom in atoms])
        return xyz.mean(axis=0)

    def copy(self) -> "LabeledMolecule":
        """Create a copy of the wrapper."""
        return LabeledMolecule(self.inner.copy(), self.label)


# Use the custom wrapper
molecule = LabeledMolecule(struct, label="test_molecule")
print(f"Label: {molecule.label}")
print(f"Number of atoms: {molecule.n_atoms}")
print(f"Center of mass: {molecule.center_of_mass}")

class LabeledMolecule(Wrapper[Atomistic]): """A custom wrapper that adds a label to a molecule.""" def __init__(self, inner: Atomistic, label: str): super().__init__(inner) self.label = label @property def n_atoms(self) -> int: """Explicitly forward the number of atoms.""" return len(self.inner.atoms) @property def center_of_mass(self): """Calculate center of mass.""" # This is a simplified example atoms = self.inner.atoms xyz = np.array([[atom["x"], atom["y"], atom["z"]] for atom in atoms]) return xyz.mean(axis=0) def copy(self) -> "LabeledMolecule": """Create a copy of the wrapper.""" return LabeledMolecule(self.inner.copy(), self.label) # Use the custom wrapper molecule = LabeledMolecule(struct, label="test_molecule") print(f"Label: {molecule.label}") print(f"Number of atoms: {molecule.n_atoms}") print(f"Center of mass: {molecule.center_of_mass}")

Label: test_molecule
Number of atoms: 3
Center of mass: [1. 1. 1.]