Compute Module¶

The compute module provides a unified abstraction for defining and executing computational operations on molecular structures.

Core Concepts¶

• Result: Base class for computation outputs
• Compute: Abstract base class for computation operations
• ComputeContext: Optional context for sharing expensive intermediates

In [1]:

from dataclasses import dataclass

from molpy.compute import Compute, ComputeContext, Result
from molpy.core.frame import Frame


# Define a custom result
@dataclass
class RadiusOfGyrationResult(Result):
    """Result from radius of gyration calculation."""

    rg: float = 0.0


# Define a custom compute
class RadiusOfGyrationCompute(Compute[Frame, RadiusOfGyrationResult]):
    """Compute radius of gyration for a frame."""

    def compute(self, input: Frame) -> RadiusOfGyrationResult:
        if "atoms" not in input:
            return RadiusOfGyrationResult(name="radius_of_gyration", rg=0.0)

        atoms = input["atoms"]
        import numpy as np

        # Extract coordinates
        x = atoms["x"][:]
        y = atoms["y"][:]
        z = atoms["z"][:]
        coords = np.column_stack([x, y, z])

        # Calculate center of mass
        com = np.mean(coords, axis=0)

        # Calculate radius of gyration
        r_squared = np.sum((coords - com) ** 2, axis=1)
        rg = np.sqrt(np.mean(r_squared))

        return RadiusOfGyrationResult(name="radius_of_gyration", rg=float(rg))


# Use the compute
compute = RadiusOfGyrationCompute()
# result = compute(frame)  # Uncomment when you have a frame

from dataclasses import dataclass from molpy.compute import Compute, ComputeContext, Result from molpy.core.frame import Frame # Define a custom result @dataclass class RadiusOfGyrationResult(Result): """Result from radius of gyration calculation.""" rg: float = 0.0 # Define a custom compute class RadiusOfGyrationCompute(Compute[Frame, RadiusOfGyrationResult]): """Compute radius of gyration for a frame.""" def compute(self, input: Frame) -> RadiusOfGyrationResult: if "atoms" not in input: return RadiusOfGyrationResult(name="radius_of_gyration", rg=0.0) atoms = input["atoms"] import numpy as np # Extract coordinates x = atoms["x"][:] y = atoms["y"][:] z = atoms["z"][:] coords = np.column_stack([x, y, z]) # Calculate center of mass com = np.mean(coords, axis=0) # Calculate radius of gyration r_squared = np.sum((coords - com) ** 2, axis=1) rg = np.sqrt(np.mean(r_squared)) return RadiusOfGyrationResult(name="radius_of_gyration", rg=float(rg)) # Use the compute compute = RadiusOfGyrationCompute() # result = compute(frame) # Uncomment when you have a frame

Sharing Context Between Computes¶

You can share expensive intermediate computations using ComputeContext:

In [2]:

# Create shared context
context = ComputeContext()
# context.data["neighbor_list"] = compute_neighbor_list(frame)

# Use in multiple computes
# compute1 = SomeCompute(context=context)
# compute2 = AnotherCompute(context=context)
# Both can access the neighbor list from context

# Create shared context context = ComputeContext() # context.data["neighbor_list"] = compute_neighbor_list(frame) # Use in multiple computes # compute1 = SomeCompute(context=context) # compute2 = AnotherCompute(context=context) # Both can access the neighbor list from context