rootstock 0.5.1
MLIP calculators with isolated Python environments
Requires Python: >=3.10

# Rootstock

A proof-of-concept package for running MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments, communicating via the i-PI protocol over Unix sockets. Currently deployed on Princeton's Della cluster.

## Quick Start

```python
from ase.build import bulk
from rootstock import RootstockCalculator

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

with RootstockCalculator(
    cluster="della",
    model="mace",
    checkpoint="medium",
    device="cuda",
) as calc:
    atoms.calc = calc
    print(atoms.get_potential_energy())
    print(atoms.get_forces())
```

**Note:** Environments must be pre-built before use. See [Administrator Setup](#administrator-setup).

## Installation

```bash
pip install rootstock
# or
uv pip install rootstock
```

## API

The `model` parameter selects the environment family. The optional `checkpoint` parameter selects specific model weights (defaults to the environment's default if omitted).

```python
# Explicit checkpoint
RootstockCalculator(cluster="della", model="mace", checkpoint="medium")

# Default checkpoint (each environment has a sensible default)
RootstockCalculator(cluster="della", model="uma")

# Custom root path instead of a known cluster
RootstockCalculator(root="/scratch/gpfs/ROSENGROUP/common/rootstock", model="mace")
```

## Available Models

| Model | Checkpoint | Description |
|-------|------------|-------------|
| `mace` | `small` | MACE-MP-0 small |
| `mace` | `medium` (default) | MACE-MP-0 medium |
| `chgnet` | _(pretrained)_ | CHGNet pretrained universal model |
| `uma` | `uma-s-1p1` (default) | Meta UMA small model (FAIRChem) |
| `tensornet` | `TensorNet-MatPES-PBE-v2025.1-PES` (default) | TensorNet MatPES PBE (MatGL) |

## Architecture

```
Main Process                          Worker Process (subprocess)
+-------------------------+          +-----------------------------+
| RootstockCalculator     |          | Pre-built venv Python       |
| (ASE-compatible)        |          | (mace_env/bin/python)       |
|                         |          |                             |
| server.py (i-PI server) |<-------->| worker.py (i-PI client)     |
| - sends positions       |   Unix   | - receives positions        |
| - receives forces       |  socket  | - calculates forces         |
+-------------------------+          +-----------------------------+
```

The worker process uses a pre-built virtual environment, providing:
- **Fast startup**: No dependency installation at runtime
- **Filesystem compatibility**: Works on NFS, Lustre, GPFS
- **Reproducibility**: Same environment every time

## Directory Structure

```
{root}/
├── .python/                # uv-managed Python interpreters (portable)
├── environments/           # Environment SOURCE files (*.py with PEP 723)
│   ├── mace_env.py
│   ├── chgnet_env.py
│   ├── uma_env.py
│   └── tensornet_env.py
├── envs/                   # Pre-built virtual environments
│   ├── mace_env/
│   │   ├── bin/python
│   │   ├── lib/python3.11/site-packages/
│   │   └── env_source.py   # Copy of environment source
│   └── ...
├── home/                   # Fake HOME for build & workers
│   ├── .cache/fairchem/    # FAIRChem weights
│   └── .matgl/             # MatGL weights
└── cache/                  # XDG_CACHE_HOME for well-behaved libs
    ├── mace/
    └── huggingface/
```

## CLI Commands

```bash
# Build a pre-built environment
rootstock build <env_name> --root <path> [--models m1,m2] [--force]

# Show status
rootstock status --root <path>

# Register an environment source file
rootstock register <env_file> --root <path>

# List environments
rootstock list --root <path>
```

## Administrator Setup

Environments must be pre-built before users can run calculations.

### 1. Create Directory Structure

```bash
mkdir -p /scratch/gpfs/ROSENGROUP/common/rootstock/{environments,envs,cache,home}
```

### 2. Register Environment Source Files

```bash
rootstock register environments/mace_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
rootstock register environments/chgnet_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
rootstock register environments/uma_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
rootstock register environments/tensornet_env.py --root /scratch/gpfs/ROSENGROUP/common/rootstock
```

### 3. Build Environments

```bash
ROOT=/scratch/gpfs/ROSENGROUP/common/rootstock

rootstock build mace_env --root $ROOT --models small,medium
rootstock build chgnet_env --root $ROOT
rootstock build uma_env --root $ROOT --models uma-s-1p1
rootstock build tensornet_env --root $ROOT

# Verify
rootstock status --root $ROOT
```

## Local Development

```bash
uv venv && source .venv/bin/activate
uv pip install -e ".[dev]"
ruff check rootstock/
ruff format rootstock/
```