Adding a Compute Operation

This page shows how to add reusable analysis operations (Compute) to MolPy.

Which base class to use

Need	Base class	Example
Reusable analysis on array data	Compute	MSD, DisplacementCorrelation
One-off calculation with no config	plain function	compute_msd(positions)

Compute is a frozen dataclass. Configuration is set at construction and cannot change. Execution goes through run() (or __call__).

Adding a Compute operation

Subclass Compute, declare configuration as frozen dataclass fields, implement run().

from dataclasses import dataclass
import numpy as np
from numpy.typing import NDArray
from molpy.compute import Compute

@dataclass(frozen=True)
class RadiusOfGyration(Compute):
    """Compute radius of gyration for each frame."""

    use_masses: bool = True

    def run(self, positions: NDArray, masses: NDArray | None = None) -> NDArray:
        """Compute Rg for a set of positions.

        Args:
            positions: shape (n_atoms, 3)
            masses: shape (n_atoms,), optional

        Returns:
            Scalar Rg value.
        """
        if self.use_masses and masses is not None:
            w = masses / masses.sum()
        else:
            w = np.ones(len(positions)) / len(positions)

        com = (positions * w[:, None]).sum(axis=0)
        dr = positions - com
        rg2 = (w * (dr ** 2).sum(axis=1)).sum()
        return float(np.sqrt(rg2))

Usage:

rg = RadiusOfGyration(use_masses=True)
value = rg(positions, masses)   # __call__ delegates to run()

Design rules

1. Configuration goes in fields — set once at init, frozen forever
2. Runtime data goes through run() — different inputs, same protocol
3. No mutation — run() returns new objects, never modifies inputs
4. Keep run() focused — one clear task, not a workflow engine
5. Test in isolation — each Compute should be testable with synthetic data

Checklist

• Subclass Compute
• Add @dataclass(frozen=True) decorator
• Implement run() with type hints
• Write tests in tests/test_compute/