dpa_tools

Lightweight fine-tuning helpers for the DPA-3.1 pretrained descriptor. Reuse the frozen DPA representation for downstream materials-property regression in a handful of lines, without standing up a full DeepMD training pipeline.

What's in the box

Three complementary workflows sit on top of the same DPA backbone:

• Path B — frozen descriptor + sklearn head. Encode each system once with the pretrained DPA descriptor, pool atomic features into a per-system vector, and train a Random Forest, Ridge, or MLP regressor on top. Seconds to fit, no GPU needed past the one-shot encode, and strong on intensive/extensive scalars where global pooling is enough.
• Training paradigms — linear probe / fine-tune / scratch. Run dp --pt train under the hood for single-task neural adaptation. LP freezes the backbone and trains a property fitting net; FT updates the full network; Scratch trains from random init. Use when Path B plateaus and you need the descriptor to move.
• MFT — multi-task fine-tuning. Jointly train a downstream property head and an auxiliary force-field head through dp --pt train, so the descriptor itself adapts while a physics-grounded aux task prevents representation collapse. Use this for OOD generalization.

All three paths share the same I/O surface (fit / predict / evaluate / freeze) and the same data layer: dpdata is the internal unique data format. Every input (file path, glob, or dpdata object) flows through load_data() and is normalized to dpdata.System / LabeledSystem — no raw .npy reads outside the data layer.

Scalar conditions (temperature, pressure, etc.) are supported via ConditionManager, which standardizes per-key condition arrays and concatenates them to descriptor features before the sklearn head.

Quickstart

Load a dataset and run cross-validation

from dpa_tools import load_dataset, DPAFineTuner
from dpa_tools.cv import cross_validate

# auto-detect system directories from a dpdata tree
systems = load_dataset("my_catalyst_data/", label_key="energy")

# frozen DPA descriptor + sklearn head, 5-fold GroupKFold
model = DPAFineTuner(
    pretrained="DPA-3.1-3M",
    pooling="mean",
    seed=42,
)
result = cross_validate(
    model, systems,
    label_key="energy",
    cv=5,
    group_by="formula",          # leak-proof grouping
    granularity="composition",   # or "frame"
)
print(f"R²: {result['aggregate']['r2_mean']:.2f}")
print(f"MAE: {result['aggregate']['mae_mean']:.2f}")

Train a frozen-DPA model and predict

from dpa_tools import DPAFineTuner

model = DPAFineTuner(
    pretrained="DPA-3.1-3M",
    pooling="mean",
)
model.fit("my_data/", target_key="energy")

# predict on new systems
result = model.predict("new_system/")
print(result.predictions)       # shape (n_frames, 1)

# evaluate against stored labels
metrics = model.evaluate("new_system/")
print(f"MAE: {metrics.mae:.4f}, R²: {metrics.r2:.4f}")

# save for later (no torch needed to reload)
model.freeze("frozen_model.pth")

Use scalar conditions (temperature, pressure, ...)

# Conditions are standardized per-key and concatenated to descriptors
model = DPAFineTuner(pretrained="DPA-3.1-3M", pooling="mean")
model.fit(
    "my_data/", target_key="energy",
    conditions={"T": T_array, "P": P_array},  # each (n_frames,)
)

# Predict with conditions — must match the keys used at fit time
result = model.predict("new_system/", conditions={"T": T_new, "P": P_new})

Freeze and deploy (no DPA checkpoint needed)

from dpa_tools import DPAPredictor

pred = DPAPredictor("frozen_model.pth")
result = pred.evaluate("my_data/")
print(f"MAE: {result.mae:.4f}, R²: {result.r2:.4f}")

Uncertainty estimation

RF has native uncertainty via per-tree prediction spread; MLP uses a committee ensemble (n_committee > 1, requires fit()). Ridge cannot produce uncertainty — it has a unique closed-form solution.

from dpa_tools import DPAPredictor

pred = DPAPredictor("frozen_model.pth")

# For RF: tree-level standard deviation (no n_committee or fit() needed)
result = pred.predict("my_data/", return_uncertainty=True)

# result.predictions  — ensemble mean, shape (n_frames, task_dim)
# result.uncertainty  — per-sample std, same shape

For MLP with committee uncertainty:

pred = DPAPredictor("frozen_model.pth", n_committee=5)
pred.fit("train_data/", target_key="energy")       # trains ensemble
result = pred.predict("new_data/", return_uncertainty=True)

The uncertainty value is the standard deviation across estimators (RF trees or MLP committee members) — larger values indicate higher model disagreement and lower confidence.

Extract descriptors for custom modeling

from dpa_tools.finetuner import extract_descriptors

X = extract_descriptors(
    "my_data/",
    pretrained="DPA-3.1-3M",
    pooling="mean+std",   # mean, sum, mean+std, mean+std+max+min
)
# X.shape == (n_frames_total, feat_dim)
# Use with any sklearn / torch model

Deterministic splits for reproducibility

from dpa_tools.cv import train_test_split, cross_validate

# fixed folds from a manifest file (paper-ready)
train, valid, test = train_test_split(
    systems, manifest="split_manifest.json",
)

# or deterministic GroupKFold (sorted groups, no shuffle)
result = cross_validate(
    model, systems, cv=5,
    group_by="formula",
    manifest="split_manifest.json",   # optional: use manifest folds
)

CLI

All Python API features are exposed through dpa-tools. No type_map flags needed for most commands — element symbols are auto-inferred from the checkpoint and data.

Cross-validate a frozen baseline (seconds)

dpa-tools cv \
    --data /path/to/dpdata \
    --label-key overpotential \
    --pretrained /path/to/DPA-3.1-3M.pt \
    --predictor rf \
    --cv 5 \
    --granularity composition

Train a single-task model (any strategy)

# linear probe — freeze backbone, train fitting net
dpa-tools fit \
    --train-data /path/to/train \
    --valid-data /path/to/valid \
    --pretrained /path/to/DPA-3.1-3M.pt \
    --strategy linear_probe \
    --property-name overpotential --task-dim 1 --intensive \
    --max-steps 3000 --output-dir ./lp_output

# finetune — full network
dpa-tools fit \
    --train-data /path/to/train --valid-data /path/to/valid \
    --pretrained /path/to/DPA-3.1-3M.pt \
    --strategy finetune \
    --property-name overpotential --task-dim 1 --intensive \
    --max-steps 3000 --output-dir ./ft_output

# frozen_sklearn — sklearn head (original Path B)
dpa-tools fit \
    --train-data /path/to/dpdata \
    --pretrained /path/to/DPA-3.1-3M.pt \
    --strategy frozen_sklearn \
    --predictor rf --target-key overpotential \
    --output frozen_model.pth

Multi-task fine-tuning

dpa-tools mft \
    --data /path/to/downstream_data \
    --aux-data /path/to/OC22 \
    --label-key overpotential \
    --pretrained /path/to/DPA-3.1-3M.pt \
    --aux-branch OC22 --aux-prob 0.5 \
    --downstream-task-type property \
    --property-name overpotential --task-dim 1 --intensive \
    --aux-batch-size auto:128 --downstream-batch-size 3 \
    --max-steps 100000 --output-dir ./mft_output

Extract descriptors for custom use

dpa-tools extract-descriptors \
    --data /path/to/systems \
    --pretrained /path/to/DPA-3.1-3M.pt \
    --pooling mean+std \
    --output descriptors.npy

Deploy a frozen bundle (no checkpoint needed)

dpa-tools predict --model frozen_model.pth --data /path/to/new_system --output pred.npy
dpa-tools evaluate --model frozen_model.pth --data /path/to/new_system

Data utilities

dpa-tools convert --input my.xyz --output ./npy_system --fmt extxyz --type-map H,O,Cu
dpa-tools batch-convert --glob "./calcs/*/OUTCAR" --output ./npy_tree --fmt vasp/outcar
dpa-tools check-data --data ./npy_system
dpa-tools attach-labels --data ./npy_system --head overpotential --values labels.npy

Data format

dpa_tools consumes the deepmd/npy directory format internally via dpdata as the unique data container. Everything else (XYZ, VASP OUTCAR/POSCAR, SDF, ...) is converted in one step with dpa_tools.convert(...), and external scalar or vector labels are attached with dpa_tools.attach_labels(...).

Internally, load_data() is the single polymorphic entry point — it normalizes file paths, glob patterns, and dpdata objects into a flat list[dpdata.System]. This means every consumer in the library (finetuner, predictor, CV, MFT) sees the same data representation, with no raw .npy reads outside the data/ subpackage.

Benchmark

All benchmarks and reproducer scripts live in the dpa-bench repository. See dpa-bench/docs/ for methodology notes per experiment.

QM9 frozen-descriptor probing (Experiment 1)

DPA-3.1-3M frozen descriptor + sklearn head (Ridge / MLP) on QM9. 10k train / 1k val / 1k test, 3 properties, 4 pooling × 2 predictor grid, ~90 min on Tesla T4.

Property	Test MAE	Test R²	Best config
HOMO–LUMO gap	0.169 eV	0.971	mean+std+max+min + mlp
Atomization energy	0.053 eV	0.996	mean + mlp
Dipole magnitude	0.463 Debye	0.827	mean+std+max+min + mlp

Frozen descriptors reach SchNet-class accuracy on small organic molecules without training the descriptor. Dipole is limited — global pooling discards directional information that equivariant readouts exploit.

QM9 BOOM multi-task fine-tuning (Experiment 2)

Reproduction of Zhang et al., "Multi-task fine-tuning" on the QM9 BOOM benchmark (split_boom). DPA-3.1-3M, 4 methods × 3 properties, seed=42, 100k steps each. ID (small molecules) + OOD (large molecules) test splits. ~36–48 hr sequentially on Tesla T4.

MAE in meV:

Method	HOMO ID	HOMO OOD	LUMO ID	LUMO OOD	GAP ID	GAP OOD
Scratch	2.15	9.70	2.43	13.04	3.40	16.04
FT	1.28	7.57	1.24	8.34	2.13	12.36
LP	5.82	15.38	8.38	17.81	10.33	24.87
MFT	1.97	10.21	2.10	8.09	3.39	14.35

MFT improves OOD generalization on LUMO (+38%) and GAP (+11%) vs Scratch, reproducing the paper's core claim. LP (linear probe / frozen descriptor) has the worst OOD performance — the descriptor needs to adapt for OOD generalization.

Pooling selection guide

Different property classes prefer different pooling strategies. Empirically, on QM9:

Property type	Recommended pooling	Predictor	Notes
Extensive scalar (energy, atomization)	sum	rf	RF with sum pooling works well out of the box
Intensive scalar (gap, HOMO, LUMO)	mean+std / mean+std+max+min	mlp	Distribution statistics carry signal
Vector magnitude (dipole, ...)	mean+std+max+min	mlp	Limited; prefer equivariant models or MFT
Local atomic properties (forces, charges)	—	—	Not supported by Path B; use MFT or training paradigms

Context embedding

For properties that depend on external scalar conditions (temperature, pressure, doping level, ...), use the conditions parameter:

model.fit(data, target_key="energy", conditions={"T": T_array})

ConditionManager fits a StandardScaler per named key and concatenates the standardized values to the pooled descriptor features. The condition keys and scalers are serialized into the frozen bundle by freeze() and restored automatically by DPAPredictor. An error is raised if conditions are expected but missing (or vice versa) at predict time.

Documentation

• docs/architecture.md — full architecture: data layer, Path B, training paradigms, MFT, pooling trade-offs, descriptor cache