Writing Environment Files¶

Each MLIP is defined by a small Python file with PEP 723 inline metadata specifying its dependencies and a setup() function that returns an ASE calculator.

Creating a New Environment¶

Use the rootstock new-env command to scaffold a new environment file from a template.

Generating a Template¶

# Create a new environment file
rootstock new-env mace

# Specify custom output path
rootstock new-env mace -o ./environments/mace_env.py

# Overwrite existing file
rootstock new-env mace --force

This generates a template file mace_env.py with the required structure:

# /// script
# requires-python = ">=3.12"
# dependencies = [
#
# ]
# ///
"""
MACE environment for Rootstock.

TODO: Add description of this environment.
"""


def setup(model: str | None = None, device: str = "cuda"):
    """
    Load a calculator.

    Args:
        model: Model identifier or checkpoint name.
        device: PyTorch device string (e.g., "cuda", "cuda:0", "cpu").

    Returns:
        ASE-compatible calculator.
    """
    raise NotImplementedError("TODO: Implement setup()")

After generating the template, fill in the dependencies and implement the setup() function as described below.

Basic Structure¶

# /// script
# requires-python = ">=3.10"
# dependencies = ["mace-torch>=0.3.14", "ase>=3.22", "torch>=2.0,<2.10"]
# ///

def setup(model: str, device: str = "cuda"):
    from mace.calculators import mace_mp
    return mace_mp(model=model, device=device, default_dtype="float32")

How It Works¶

PEP 723 metadata: The # /// script block defines Python version requirements and dependencies
setup() function: Called once when a worker starts. The returned calculator is reused for all calculations
Environment building: Rootstock uses uv to create an isolated virtual environment from the specified dependencies

Required Elements¶

PEP 723 Metadata Block¶

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "mace-torch>=0.3.14",
#     "ase>=3.22",
#     "torch>=2.0,<2.10"
# ]
# ///

requires-python: Minimum Python version
dependencies: List of pip-installable packages with version constraints

The `setup()` Function¶

def setup(model: str, device: str = "cuda"):
    # Import the MLIP library
    from mace.calculators import mace_mp

    # Create and return an ASE calculator
    return mace_mp(model=model, device=device, default_dtype="float32")

Parameters:

model (str): Checkpoint or model name passed by the user
device (str): PyTorch device ("cuda" or "cpu")

Returns:

An ASE-compatible calculator object.

Examples¶

MACE¶

# /// script
# requires-python = ">=3.10"
# dependencies = ["mace-torch>=0.3.14", "ase>=3.22", "torch>=2.0,<2.10"]
# ///

def setup(model: str, device: str = "cuda"):
    from mace.calculators import mace_mp
    return mace_mp(model=model, device=device, default_dtype="float32")

CHGNet¶

# /// script
# requires-python = ">=3.10"
# dependencies = ["chgnet>=0.3.0", "ase>=3.22", "torch>=2.0"]
# ///

def setup(model: str = "pretrained", device: str = "cuda"):
    from chgnet.model import CHGNetCalculator
    return CHGNetCalculator(use_device=device)

UMA (FAIRChem)¶

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "fairchem-core>=1.0.0",
#     "ase>=3.22",
#     "torch>=2.0"
# ]
# ///

def setup(model: str = "uma-s-1p1", device: str = "cuda"):
    from fairchem.core import OCPCalculator
    return OCPCalculator(checkpoint_path=model, device=device)

TensorNet (MatGL)¶

# /// script
# requires-python = ">=3.10"
# dependencies = ["matgl>=1.0.0", "ase>=3.22", "torch>=2.0"]
# ///

def setup(model: str = "TensorNet-MatPES-PBE-v2025.1-PES", device: str = "cuda"):
    import matgl
    from matgl.ext.ase import MatGLCalculator

    potential = matgl.load_model(model)
    return MatGLCalculator(potential, device=device)

Best Practices¶

Pin dependency versions¶

Use version constraints to ensure reproducible builds:

# Good: pinned versions
# dependencies = ["mace-torch>=0.3.14,<0.4", "torch>=2.0,<2.10"]

# Avoid: unpinned versions
# dependencies = ["mace-torch", "torch"]

Handle model loading errors gracefully¶

def setup(model: str, device: str = "cuda"):
    from mace.calculators import mace_mp

    try:
        return mace_mp(model=model, device=device, default_dtype="float32")
    except Exception as e:
        raise RuntimeError(f"Failed to load MACE model '{model}': {e}")

Document available checkpoints¶

Add a comment listing known-good checkpoints:

# /// script
# requires-python = ">=3.10"
# dependencies = ["mace-torch>=0.3.14", "ase>=3.22", "torch>=2.0,<2.10"]
# ///
#
# Available checkpoints: small, medium, large

def setup(model: str, device: str = "cuda"):
    from mace.calculators import mace_mp
    return mace_mp(model=model, device=device, default_dtype="float32")

Testing Your Environment¶

After creating an environment file, install it and verify the build:

# Install the environment
rootstock install my_env.py --models default_checkpoint

# Verify it was built successfully
rootstock status

To test the calculator, create a simple Python script:

from ase.build import bulk
from rootstock import RootstockCalculator

atoms = bulk("Cu", "fcc", a=3.6)

with RootstockCalculator(
    root="/path/to/rootstock",
    model="my_env",
    checkpoint="default_checkpoint",
    device="cuda",
) as calc:
    atoms.calc = calc
    energy = atoms.get_potential_energy()
    forces = atoms.get_forces()
    print(f"Energy: {energy:.4f} eV")
    print(f"Forces shape: {forces.shape}")