Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,11 @@
+outputs/
+checkpoints/
+datasets/
+weight/esen_*.pt
+weight/esen_*.pth
+__pycache__/
+demo/__pycache__/
+*.py[cod]
+*.log
+*.traj
+relaxed.cif

README.md

@@ -0,0 +1,234 @@
+---
+datasets:
+  - OneScience-Group/oxides
+frameworks:
+  - ""
+language:
+  - en
+license: mit
+tags:
+  - OneScience
+  - eSEN
+  - materials-science
+  - machine-learning-potential
+  - equivariant-graph-neural-network
+  - structure-relaxation
+  - molecular-dynamics
+  - fine-tuning
+tasks: []
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">eSEN</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+eSEN (equivariant Smooth Energy Network) is an equivariant graph neural network interatomic potential proposed by FAIR Chemistry. It predicts the energy, atomic forces, and stress of material structures. OneScience provides an ASE calculator and fine-tuning entry points for single-point calculations, structure relaxation, molecular dynamics, and fine-tuning on custom materials data.
+
+Upstream implementation: [FAIR-Chem/fairchem](https://github.com/FAIR-Chem/fairchem)
+
+# Model Description
+
+eSEN takes periodic atomic structure graphs as input, learns a smooth potential energy surface through rotationally equivariant representations, and obtains conservative forces from energy gradients. Different pretrained checkpoints correspond to different materials domains. Select weights whose target elemental systems and DFT labeling settings are close to your use case.
+
+This repository contains the eSEN model code, inference scripts, an oxide PBE fine-tuning example, and multi-device launch configurations. Restricted pretrained eSEN checkpoints are not distributed with this repository.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :--- |
+| Single-point calculation | Predict the total energy, atomic forces, and stress of a periodic structure |
+| Structure relaxation | Optimize atomic positions and, optionally, the unit cell with ASE BFGS |
+| Molecular dynamics | Run NVT trajectories with ASE Langevin dynamics |
+| Oxide PBE fine-tuning | Fine-tune an MPTrj checkpoint using energy and force labels |
+| Distributed fine-tuning | Support single-node multi-device and multi-node Slurm DDP |
+
+This repository does not provide a complete pretraining workflow from random initialization. Its training entry point is intended for checkpoint fine-tuning.
+
+# Usage
+
+## 1. Using OneCode
+
+Try intelligent, one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware requirements**
+
+- A GPU or DCU is recommended for inference and fine-tuning.
+- A CPU can be used for import and configuration checks but is not recommended for production workloads.
+- A DCU requires a DTK runtime matching the PyTorch build.
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Group/eSEN --local-dir ./eSEN
+cd eSEN
+```
+
+### Install the Runtime Environment
+
+**DCU environment**
+
+```bash
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[matchem-dcu] \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+**GPU environment**
+
+```bash
+conda create -n onescience311 python=3.11 -y \
+  libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+pip install onescience[matchem-gpu] \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+### Pretrained Weights
+
+Access to pretrained eSEN checkpoints must be requested from the official FAIR Chemistry model release page. This repository neither provides nor redistributes those weights. After obtaining access, place the required checkpoint at:
+
+```text
+weight/
+├── Jd.pt                        # Rotation-basis file included in this repository
+├── esen_30m_mptrj.pt            # Add after obtaining access
+├── esen_30m_omat.pt             # Add after obtaining access
+└── esen_30m_oam.pt              # Add after obtaining access
+```
+
+Checkpoint details:
+
+| Checkpoint | Training domain | Recommended use |
+| --- | --- | --- |
+| `esen_30m_mptrj.pt` | MPTrj | PBE/PBE+U inference for inorganic crystals and fine-tuning in a similar labeling domain |
+| `esen_30m_omat.pt` | OMat24 | A broader range of non-equilibrium inorganic structures |
+| `esen_30m_oam.pt` | OAM | General-purpose materials pretraining starting point |
+
+Access and download resources:
+
+- FAIR Chemistry model page: https://huggingface.co/fairchem
+- Official FAIR Chemistry repository: https://github.com/FAIR-Chem/fairchem
+
+`Jd.pt` is included in the repository under `weight/`. The inference and fine-tuning scripts use it automatically, with no additional download or UMA dependency required.
+
+### Fine-Tuning Dataset
+
+This repository does not include the training data. The oxide PBE example dataset is published separately on Hugging Face:
+
+```bash
+hf download --dataset OneScience-Group/oxides \
+  --local-dir ./datasets/oxides
+```
+
+The downloaded data has the following layout:
+
+```text
+datasets/oxides/data/OXIDES/prepared/
+├── train.db
+├── val.db
+└── test.db
+```
+
+The data has already been converted to ASE DB format. The default configuration uses energy and force supervision and does not train on stress. Use `prepare_oxide_dataset.py` to regenerate the data from the official oxide JSON files.
+
+### Inference
+
+Predict single-point energy, forces, and stress:
+
+```bash
+python single_point.py --checkpoint weight/esen_30m_mptrj.pt
+```
+
+Relax a structure:
+
+```bash
+python relax.py \
+  --checkpoint weight/esen_30m_mptrj.pt \
+  --fmax 0.05 --steps 100 --output relaxed.cif
+```
+
+Run NVT molecular dynamics:
+
+```bash
+python md.py \
+  --checkpoint weight/esen_30m_mptrj.pt \
+  --steps 100 --temperature 300 --timestep 1.0 --output md.traj
+```
+
+Use `--input` to read CIF, POSCAR, XYZ, or any other structure format supported by ASE.
+
+### Fine-Tuning
+
+Fine-tune on oxide PBE data with one device:
+
+```bash
+bash demo/run.sh --config configs/finetune_1dcu.yaml
+```
+
+Main configuration fields:
+
+| YAML field | Purpose |
+| --- | --- |
+| `checkpoint` | Path to the initialization checkpoint |
+| `train`, `val` | Paths to ASE DB or ASE-LMDB data |
+| `epochs`, `batch_size`, `workers` | Number of training epochs, batch size, and data-loading workers |
+| `lr` | AdamW learning rate |
+| `energy_weight`, `force_weight`, `stress_weight` | Loss weight for each supervised target |
+| `fit_element_references` | Whether to refit elemental energy references on the training set |
+| `launch.num_nodes`, `launch.num_gpus` | Number of nodes and devices per node |
+| `launch.mode` | Run directly with `local` or submit to Slurm with `submit` |
+| `slurm.*` | Slurm partition, time limit, and CPU resources |
+
+Use the same entry point for multi-device and multi-node jobs by selecting another YAML file:
+
+```bash
+bash demo/run.sh --config configs/finetune_2dcu.yaml
+bash demo/run.sh --config configs/finetune_16dcu.yaml
+```
+
+When the current node has fewer visible devices than requested, or when the configuration requests multiple nodes, `demo/run.sh` automatically submits the job through `sbatch`.
+
+Evaluate a fine-tuned checkpoint:
+
+```bash
+python evaluate.py \
+  --checkpoint outputs/<run>/checkpoints/esen_oxides_1dcu_finetuned.pt \
+  --data datasets/oxides/data/OXIDES/prepared/test.db
+```
+
+# File Reference
+
+| Path | Purpose |
+| --- | --- |
+| `model/` | Source code for the eSEN backbone, prediction heads, and graph interface |
+| `single_point.py` | Single-point energy, force, and stress inference |
+| `relax.py` | Periodic structure relaxation |
+| `md.py` | NVT molecular dynamics |
+| `finetune.py` | Low-level checkpoint fine-tuning entry point |
+| `evaluate.py` | Error evaluation on an independent dataset |
+| `prepare_oxide_dataset.py` | Convert oxide JSON data to ASE DB format |
+| `demo/run.sh` | YAML-driven local and Slurm fine-tuning entry point |
+| `demo/configs/` | Single-device, multi-device, and multi-node fine-tuning configurations |
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- The eSEN model code is adapted from FairChem Core and follows the upstream FairChem MIT License. The model code and the included `Jd.pt` file are used under that model license.
+- eSEN checkpoints are not distributed with this repository. Follow the access conditions and licenses specified on the FAIR Chemistry model pages.
+- When using OMat24, MPTrj, OAM, or oxide data, cite the datasets and corresponding model work actually used.

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

demo/_parse_config.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+"""Extract eSEN demo metadata and the native finetune configuration from YAML."""
+
+from __future__ import annotations
+
+import sys
+
+import yaml
+
+
+META_KEYS = {"name", "description", "launch", "slurm", "nccl"}
+
+
+def _config(path: str) -> dict:
+    with open(path, encoding="utf-8") as handle:
+        return yaml.safe_load(handle) or {}
+
+
+def _launch(cfg: dict) -> dict:
+    return cfg.get("launch", {}) or {}
+
+
+def _slurm(cfg: dict) -> dict:
+    return cfg.get("slurm", {}) or {}
+
+
+def _print_launch(cfg: dict) -> None:
+    launch = _launch(cfg)
+    mode = launch.get("mode", "local")
+    if mode not in {"local", "submit"}:
+        raise ValueError("launch.mode must be 'local' or 'submit'")
+    print(f"RUN_MODE={mode}")
+    print(f"NODES={launch.get('num_nodes', 1)}")
+    print(f"GPUS_PER_NODE={launch.get('num_gpus', 1)}")
+    print(f"OMP_NUM_THREADS={launch.get('omp_num_threads', 1)}")
+
+
+def _print_slurm(cfg: dict) -> None:
+    slurm = _slurm(cfg)
+    print(f"PARTITION={slurm.get('partition', 'hx1hdexclu12')}")
+    print(f"TIME={slurm.get('time', '24:00:00')}")
+    print(f"CPUS_PER_TASK={slurm.get('cpus_per_task', 16)}")
+
+
+def _print_env(cfg: dict) -> None:
+    nccl = cfg.get("nccl", {}) or {}
+    if nccl.get("socket_ifname"):
+        print(f"export NCCL_SOCKET_IFNAME={nccl['socket_ifname']}")
+    if nccl.get("ib_hca"):
+        print(f"export NCCL_IB_HCA={nccl['ib_hca']}")
+    if nccl.get("proto"):
+        print(f"export NCCL_PROTO={nccl['proto']}")
+    print("export HSA_FORCE_FINE_GRAIN_PCIE=1")
+
+
+def _print_finetune_config(cfg: dict) -> None:
+    finetune = {key: value for key, value in cfg.items() if key not in META_KEYS}
+    yaml.safe_dump(
+        finetune,
+        sys.stdout,
+        sort_keys=False,
+        default_flow_style=False,
+        allow_unicode=True,
+    )
+
+
+def main() -> None:
+    if len(sys.argv) != 3:
+        raise SystemExit(
+            "usage: _parse_config.py <config.yaml> <name|launch|slurm|env|finetune-config>"
+        )
+    cfg = _config(sys.argv[1])
+    action = sys.argv[2]
+    actions = {
+        "name": lambda: print(cfg.get("name", "esen_finetune")),
+        "launch": lambda: _print_launch(cfg),
+        "slurm": lambda: _print_slurm(cfg),
+        "env": lambda: _print_env(cfg),
+        "finetune-config": lambda: _print_finetune_config(cfg),
+    }
+    try:
+        actions[action]()
+    except KeyError as error:
+        raise SystemExit(f"unknown action: {action}") from error
+
+
+if __name__ == "__main__":
+    main()

demo/configs/finetune_16dcu.yaml

@@ -0,0 +1,36 @@
+# eSEN oxide PBE fine-tuning: 2 nodes x 8 DCUs.
+# Submit from examples/matchem/esen:
+#   bash demo/run.sh --config configs/finetune_16dcu.yaml
+name: esen_oxides_16dcu
+description: "eSEN oxide PBE fine-tuning on 2 nodes x 8 DCUs"
+
+launch:
+  mode: submit
+  num_nodes: 2
+  num_gpus: 8
+  omp_num_threads: 1
+
+slurm:
+  partition: hx1hdexclu12
+  time: "24:00:00"
+  cpus_per_task: 128
+
+# Set fabric-specific NCCL values here when required by the cluster.
+nccl: {}
+
+checkpoint: "${ESEN_REPO_DIR}/weight/esen_30m_mptrj.pt"
+train: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/train.db"
+val: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/val.db"
+output: "checkpoints/esen_oxides_16dcu_finetuned.pt"
+
+device: cuda
+backend: nccl
+seed: 0
+fit_element_references: true
+epochs: 10
+batch_size: 2
+workers: 0
+lr: 1.0e-6
+energy_weight: 1.0
+force_weight: 10.0
+stress_weight: 0.0

demo/configs/finetune_16dcu_smoke.yaml

@@ -0,0 +1,36 @@
+# End-to-end validation only: 2 nodes x 8 DCUs, one batch per rank.
+# This is not a scientific fine-tuning recipe.
+name: esen_oxides_16dcu_smoke
+description: "eSEN two-node DDP smoke test"
+
+launch:
+  mode: submit
+  num_nodes: 2
+  num_gpus: 8
+  omp_num_threads: 1
+
+slurm:
+  partition: hx1hdexclu12
+  time: "00:20:00"
+  cpus_per_task: 128
+
+nccl: {}
+
+checkpoint: "${ESEN_REPO_DIR}/weight/esen_30m_mptrj.pt"
+train: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/train.db"
+val: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/val.db"
+output: "checkpoints/esen_oxides_16dcu_smoke.pt"
+
+device: cuda
+backend: nccl
+seed: 0
+fit_element_references: true
+epochs: 1
+batch_size: 2
+workers: 0
+max_train_samples: 16
+max_val_samples: 16
+lr: 1.0e-6
+energy_weight: 1.0
+force_weight: 10.0
+stress_weight: 0.0

demo/configs/finetune_1dcu.yaml

@@ -0,0 +1,35 @@
+# eSEN oxide PBE fine-tuning: 1 node x 1 DCU.
+# Run from examples/matchem/esen:
+#   bash demo/run.sh --config configs/finetune_1dcu.yaml
+name: esen_oxides_1dcu
+description: "eSEN oxide PBE fine-tuning on one DCU"
+
+launch:
+  mode: local
+  num_nodes: 1
+  num_gpus: 1
+  omp_num_threads: 1
+
+slurm:
+  partition: hx1hdexclu12
+  time: "24:00:00"
+  cpus_per_task: 16
+
+nccl: {}
+
+checkpoint: "${ESEN_REPO_DIR}/weight/esen_30m_mptrj.pt"
+train: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/train.db"
+val: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/val.db"
+output: "checkpoints/esen_oxides_1dcu_finetuned.pt"
+
+device: cuda
+backend: nccl
+seed: 0
+fit_element_references: true
+epochs: 10
+batch_size: 2
+workers: 0
+lr: 1.0e-6
+energy_weight: 1.0
+force_weight: 10.0
+stress_weight: 0.0

demo/configs/finetune_2dcu.yaml

@@ -0,0 +1,35 @@
+# eSEN oxide PBE fine-tuning: 1 node x 2 DCUs.
+# Run from examples/matchem/esen:
+#   bash demo/run.sh --config configs/finetune_2dcu.yaml
+name: esen_oxides_2dcu
+description: "eSEN oxide PBE fine-tuning on 1 node x 2 DCUs"
+
+launch:
+  mode: local
+  num_nodes: 1
+  num_gpus: 2
+  omp_num_threads: 1
+
+slurm:
+  partition: hx1hdexclu12
+  time: "24:00:00"
+  cpus_per_task: 32
+
+nccl: {}
+
+checkpoint: "${ESEN_REPO_DIR}/weight/esen_30m_mptrj.pt"
+train: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/train.db"
+val: "${ESEN_DATASET_DIR}/data/OXIDES/prepared/val.db"
+output: "checkpoints/esen_oxides_2dcu_finetuned.pt"
+
+device: cuda
+backend: nccl
+seed: 0
+fit_element_references: true
+epochs: 10
+batch_size: 2
+workers: 0
+lr: 1.0e-6
+energy_weight: 1.0
+force_weight: 10.0
+stress_weight: 0.0

demo/run.sh

@@ -0,0 +1,147 @@
+#!/bin/bash
+# Run native eSEN fine-tuning directly or submit it to Slurm.
+set -euo pipefail
+
+DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+ESEN_DIR="$(cd "$DEMO_DIR/.." && pwd)"
+PARSER="$DEMO_DIR/_parse_config.py"
+CONFIG=""
+SUBMIT=false
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --config) CONFIG="$2"; shift 2 ;;
+        --config=*) CONFIG="${1#*=}"; shift ;;
+        --submit) SUBMIT=true; shift ;;
+        -h|--help)
+            echo "Usage: bash demo/run.sh --config configs/<name>.yaml"
+            echo "Set launch.mode: submit in YAML to submit a Slurm job."
+            exit 0
+            ;;
+        *) echo "Unknown argument: $1" >&2; exit 2 ;;
+    esac
+done
+
+[[ -n "$CONFIG" ]] || { echo "Please specify --config configs/<name>.yaml" >&2; exit 2; }
+[[ "$CONFIG" = /* ]] || CONFIG="$DEMO_DIR/$CONFIG"
+[[ -f "$CONFIG" ]] || { echo "Config not found: $CONFIG" >&2; exit 2; }
+
+if [[ -z "${CONDA_PREFIX:-}" ]]; then
+    echo "Activate a OneScience MatChem conda environment before running this script." >&2
+    exit 2
+fi
+
+export ESEN_REPO_DIR="$ESEN_DIR"
+export ESEN_DATASET_DIR="${ESEN_DATASET_DIR:-$ESEN_DIR/datasets/oxides}"
+export ONESCIENCE_ESEN_JD_PATH="${ONESCIENCE_ESEN_JD_PATH:-$ESEN_DIR/weight/Jd.pt}"
+
+NAME="$(python3 "$PARSER" "$CONFIG" name)"
+eval "$(python3 "$PARSER" "$CONFIG" launch)"
+eval "$(python3 "$PARSER" "$CONFIG" slurm)"
+ENV_EXPORTS="$(python3 "$PARSER" "$CONFIG" env)"
+
+if [[ "$RUN_MODE" == "submit" ]]; then
+    SUBMIT=true
+fi
+
+if ! $SUBMIT; then
+    AUTO_SUBMIT_REASON=""
+    if (( NODES > 1 )); then
+        AUTO_SUBMIT_REASON="the config requests $NODES nodes"
+    else
+        AVAILABLE_GPUS="$(python3 -c 'import torch; print(torch.cuda.device_count() if torch.cuda.is_available() else 0)' 2>/dev/null || echo 0)"
+        if (( AVAILABLE_GPUS < GPUS_PER_NODE )); then
+            AUTO_SUBMIT_REASON="the config requests $GPUS_PER_NODE DCUs but only $AVAILABLE_GPUS are visible"
+        fi
+    fi
+    if [[ -n "$AUTO_SUBMIT_REASON" ]]; then
+        if ! command -v sbatch >/dev/null 2>&1; then
+            echo "Cannot run locally: $AUTO_SUBMIT_REASON, and sbatch is unavailable." >&2
+            exit 2
+        fi
+        echo "Local resources are insufficient: $AUTO_SUBMIT_REASON. Submitting to Slurm."
+        SUBMIT=true
+    fi
+fi
+
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+OUTPUT_ROOT="${ESEN_OUTPUT_ROOT:-$ESEN_DIR/outputs}"
+OUTPUT_DIR="$OUTPUT_ROOT/${NAME}_${TIMESTAMP}"
+mkdir -p "$OUTPUT_DIR/checkpoints"
+cp "$CONFIG" "$OUTPUT_DIR/config.yaml"
+FINETUNE_CONFIG="$OUTPUT_DIR/finetune.yaml"
+python3 "$PARSER" "$CONFIG" finetune-config > "$FINETUNE_CONFIG"
+
+if $SUBMIT; then
+    SLURM_SCRIPT="$OUTPUT_DIR/submit.sh"
+    cat > "$SLURM_SCRIPT" <<EOF
+#!/bin/bash
+#SBATCH --job-name=$NAME
+#SBATCH --partition=$PARTITION
+#SBATCH --nodes=$NODES
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=$CPUS_PER_TASK
+#SBATCH --gres=dcu:$GPUS_PER_NODE
+#SBATCH --time=$TIME
+#SBATCH --output=$OUTPUT_DIR/slurm_%j.out
+#SBATCH --error=$OUTPUT_DIR/slurm_%j.err
+EOF
+
+    cat >> "$SLURM_SCRIPT" <<EOF
+
+set -euo pipefail
+export ESEN_REPO_DIR="$ESEN_DIR"
+export ESEN_DATASET_DIR="$ESEN_DATASET_DIR"
+export ONESCIENCE_ESEN_JD_PATH="$ONESCIENCE_ESEN_JD_PATH"
+export OMP_NUM_THREADS="$OMP_NUM_THREADS"
+# A login/compute shell may carry a single-device selection into sbatch.
+# Let Slurm expose all devices requested by this YAML.
+if (( $GPUS_PER_NODE > 1 )); then
+    unset HIP_VISIBLE_DEVICES
+fi
+$ENV_EXPORTS
+cd "$OUTPUT_DIR"
+EOF
+
+    if (( NODES > 1 )); then
+        cat >> "$SLURM_SCRIPT" <<EOF
+
+nodes=(\$(scontrol show hostnames "\$SLURM_JOB_NODELIST"))
+export MASTER_ADDR="\${nodes[0]}"
+export MASTER_PORT=\$((20000 + SLURM_JOB_ID % 20000))
+echo "eSEN multi-node DDP: nodes=\$SLURM_NNODES, devices/node=$GPUS_PER_NODE"
+echo "MASTER_ADDR=\$MASTER_ADDR MASTER_PORT=\$MASTER_PORT"
+
+srun --nodes="\$SLURM_NNODES" --ntasks="\$SLURM_NNODES" --ntasks-per-node=1 \\
+  bash -c 'exec torchrun \\
+    --nnodes=$NODES \\
+    --node_rank="\${SLURM_NODEID}" \\
+    --nproc_per_node=$GPUS_PER_NODE \\
+    --rdzv_id="\${SLURM_JOB_ID}" \\
+    --rdzv_backend=c10d \\
+    --rdzv_endpoint="\${MASTER_ADDR}:\${MASTER_PORT}" \\
+    "$ESEN_DIR/finetune.py" --config "$FINETUNE_CONFIG"'
+EOF
+    elif (( GPUS_PER_NODE > 1 )); then
+        cat >> "$SLURM_SCRIPT" <<EOF
+
+torchrun --standalone --nproc_per_node=$GPUS_PER_NODE \\
+  "$ESEN_DIR/finetune.py" --config "$FINETUNE_CONFIG"
+EOF
+    else
+        echo "python \"$ESEN_DIR/finetune.py\" --config \"$FINETUNE_CONFIG\"" >> "$SLURM_SCRIPT"
+    fi
+
+    chmod u+x "$SLURM_SCRIPT"
+    echo "Submitting eSEN job: $SLURM_SCRIPT"
+    sbatch "$SLURM_SCRIPT"
+    exit 0
+fi
+
+eval "$ENV_EXPORTS"
+cd "$OUTPUT_DIR"
+if (( GPUS_PER_NODE > 1 )); then
+    exec torchrun --standalone --nproc_per_node="$GPUS_PER_NODE" \
+        "$ESEN_DIR/finetune.py" --config "$FINETUNE_CONFIG"
+fi
+exec python "$ESEN_DIR/finetune.py" --config "$FINETUNE_CONFIG"

evaluate.py

@@ -0,0 +1,100 @@
+"""Evaluate an eSEN checkpoint in physical units on an independent ASE DB."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_ESEN_JD_PATH",
+    os.path.join(os.path.dirname(__file__), "weight", "Jd.pt"),
+)
+
+import torch
+
+from onescience.utils.esen.checkpoint import ESENCheckpointTransforms
+from onescience.utils.uma.common.utils import load_model_and_weights_from_checkpoint
+
+from finetune import _loader
+
+
+def _stats(error: torch.Tensor) -> dict[str, float]:
+    error = error.detach().reshape(-1).double()
+    return {
+        "mae": float(error.abs().mean()),
+        "rmse": float(error.square().mean().sqrt()),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--data", required=True, help="Independent ASE DB/ASE-LMDB")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--workers", type=int, default=0)
+    parser.add_argument("--max-samples", type=int)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--output")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA/DCU was requested but is unavailable")
+    import onescience.models.esen  # noqa: F401
+
+    model = load_model_and_weights_from_checkpoint(args.checkpoint).to(device)
+    model.eval()
+    transforms = ESENCheckpointTransforms.from_checkpoint(args.checkpoint).to(device)
+    loader = _loader(
+        args.data,
+        args.batch_size,
+        args.workers,
+        max_samples=args.max_samples,
+        train=False,
+        seed=args.seed,
+    )
+
+    energy_errors = []
+    energy_per_atom_errors = []
+    force_errors = []
+    stress_errors = []
+    with torch.enable_grad():
+        for batch in loader:
+            batch = batch.to(device)
+            prediction = model(batch)
+            pred_energy = transforms.denormalize_prediction("energy", prediction["energy"], batch)
+            target_energy = batch.energy.reshape_as(pred_energy)
+            energy_errors.append((pred_energy - target_energy).detach().cpu())
+            natoms = batch.natoms.to(pred_energy).reshape((-1,) + (1,) * (pred_energy.ndim - 1))
+            energy_per_atom_errors.append(((pred_energy - target_energy) / natoms).detach().cpu())
+
+            pred_forces = transforms.denormalize_prediction("forces", prediction["forces"], batch)
+            target_forces = batch.forces.reshape_as(pred_forces)
+            force_errors.append((pred_forces - target_forces).detach().cpu())
+            if "stress" in prediction and hasattr(batch, "stress"):
+                pred_stress = transforms.denormalize_prediction("stress", prediction["stress"], batch)
+                target_stress = batch.stress.reshape_as(pred_stress)
+                stress_errors.append((pred_stress - target_stress).detach().cpu())
+
+    result = {
+        "checkpoint": str(Path(args.checkpoint).expanduser()),
+        "data": str(Path(args.data).expanduser()),
+        "samples": len(loader.dataset),
+        "energy_total_eV": _stats(torch.cat(energy_errors)),
+        "energy_per_atom_eV": _stats(torch.cat(energy_per_atom_errors)),
+        "forces_eV_per_A": _stats(torch.cat(force_errors)),
+    }
+    if stress_errors:
+        result["stress_eV_per_A3"] = _stats(torch.cat(stress_errors))
+    print(json.dumps(result, indent=2, sort_keys=True))
+    if args.output:
+        output = Path(args.output).expanduser()
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
+
+
+if __name__ == "__main__":
+    main()

finetune.py

@@ -0,0 +1,372 @@
+"""Fine-tune an eSEN checkpoint on an ASE database.
+
+The input database must contain ASE calculator results (energy, forces, and
+optionally stress). The output checkpoint keeps the native OneScience model
+configuration and can be loaded by ``eSENCalculator.from_checkpoint``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_ESEN_JD_PATH",
+    os.path.join(os.path.dirname(__file__), "weight", "Jd.pt"),
+)
+
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Subset
+from torch.utils.data.distributed import DistributedSampler
+
+from onescience.datapipes.materials.custom_stack import data_list_collater
+from onescience.datapipes.materials.custom_stack.storage.ase_datasets import AseDBDataset
+from onescience.utils.esen.checkpoint import ESENCheckpointTransforms
+from onescience.utils.uma.normalization.element_references import (
+    fit_linear_references,
+)
+from onescience.utils.uma.common.utils import load_model_and_weights_from_checkpoint
+
+
+@dataclass(frozen=True)
+class DistributedContext:
+    """Runtime information for a normal Python process or a torchrun worker."""
+
+    rank: int = 0
+    world_size: int = 1
+    local_rank: int = 0
+
+    @property
+    def enabled(self) -> bool:
+        return self.world_size > 1
+
+    @property
+    def is_main(self) -> bool:
+        return self.rank == 0
+
+
+def _init_distributed(device_name: str, backend: str) -> DistributedContext:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size == 1:
+        if device_name.startswith("cuda") and torch.cuda.is_available():
+            torch.cuda.set_device(0)
+        return DistributedContext()
+
+    if not torch.distributed.is_available():
+        raise RuntimeError("torch.distributed is required for multi-device fine-tuning.")
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    if device_name.startswith("cuda"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("torchrun requested multiple CUDA/DCU devices, but CUDA is unavailable.")
+        torch.cuda.set_device(local_rank)
+    torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size)
+    return DistributedContext(rank=rank, world_size=world_size, local_rank=local_rank)
+
+
+def _close_distributed(context: DistributedContext) -> None:
+    if context.enabled and torch.distributed.is_initialized():
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
+def _loader(
+    path: str | list[str],
+    batch_size: int,
+    workers: int,
+    max_samples: int | None = None,
+    context: DistributedContext | None = None,
+    train: bool = False,
+    seed: int = 0,
+) -> DataLoader:
+    dataset = AseDBDataset(
+        {
+            "src": path,
+            "a2g_args": {
+                "r_edges": False,
+                "r_energy": True,
+                "r_forces": True,
+                "r_stress": True,
+            },
+        }
+    )
+    if max_samples is not None:
+        sample_count = min(max_samples, len(dataset))
+        generator = torch.Generator().manual_seed(seed)
+        indices = torch.randperm(len(dataset), generator=generator)[:sample_count].tolist()
+        dataset = Subset(dataset, indices)
+    context = context or DistributedContext()
+    sampler = None
+    if context.enabled:
+        sampler = DistributedSampler(
+            dataset,
+            num_replicas=context.world_size,
+            rank=context.rank,
+            shuffle=train,
+            drop_last=False,
+        )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=sampler is None and train,
+        sampler=sampler,
+        num_workers=workers,
+        collate_fn=lambda items: data_list_collater(items, otf_graph=True),
+    )
+
+
+def _loss(
+    pred,
+    batch,
+    energy_weight: float,
+    force_weight: float,
+    stress_weight: float,
+    transforms: ESENCheckpointTransforms,
+):
+    losses = {}
+    if energy_weight:
+        energy_target = transforms.normalize_target(
+            "energy", batch.energy, pred["energy"], batch
+        )
+        energy_error = pred["energy"] - energy_target
+        natoms_shape = (-1,) + (1,) * (energy_error.ndim - 1)
+        natoms = batch.natoms.to(energy_error).reshape(natoms_shape)
+        losses["energy"] = (energy_error / natoms).square().mean()
+    if force_weight:
+        force_target = transforms.normalize_target(
+            "forces", batch.forces, pred["forces"], batch
+        )
+        losses["forces"] = (pred["forces"] - force_target).square().mean()
+    if stress_weight and hasattr(batch, "stress") and "stress" in pred:
+        stress_target = transforms.normalize_target(
+            "stress", batch.stress, pred["stress"], batch
+        )
+        losses["stress"] = (pred["stress"] - stress_target).square().mean()
+    total = energy_weight * losses.get("energy", 0.0)
+    total = total + force_weight * losses.get("forces", 0.0)
+    total = total + stress_weight * losses.get("stress", 0.0)
+    return total, {key: float(value.detach()) for key, value in losses.items()}
+
+
+def _run_epoch(
+    model,
+    loader,
+    device,
+    optimizer,
+    weights,
+    transforms: ESENCheckpointTransforms,
+    context: DistributedContext,
+):
+    training = optimizer is not None
+    model.train(training)
+    total = 0.0
+    batches = 0
+    metric_names = tuple(
+        name
+        for name, weight in zip(("energy", "forces", "stress"), weights)
+        if weight
+    )
+    metrics = {name: 0.0 for name in metric_names}
+    for batch in loader:
+        batch = batch.to(device)
+        if training:
+            optimizer.zero_grad(set_to_none=True)
+        prediction = model(batch)
+        loss, batch_metrics = _loss(prediction, batch, *weights, transforms)
+        if training:
+            loss.backward()
+            optimizer.step()
+        total += float(loss.detach())
+        batches += 1
+        for key, value in batch_metrics.items():
+            metrics[key] = metrics.get(key, 0.0) + value
+    if batches == 0:
+        raise RuntimeError("The dataset contains no samples.")
+    values = torch.tensor([total, *metrics.values(), float(batches)], dtype=torch.float64, device=device)
+    if context.enabled:
+        torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.SUM)
+    global_batches = values[-1].item()
+    return {
+        "loss": values[0].item() / global_batches,
+        **{
+            key: values[index].item() / global_batches
+            for index, key in enumerate(metrics, start=1)
+        },
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", help="YAML configuration path")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--train", help="ASE DB or ASE-LMDB training path")
+    parser.add_argument("--val", help="ASE DB or ASE-LMDB validation path")
+    parser.add_argument("--output")
+    parser.add_argument("--device")
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--batch-size", type=int)
+    parser.add_argument("--workers", type=int)
+    parser.add_argument("--lr", type=float)
+    parser.add_argument("--energy-weight", type=float)
+    parser.add_argument("--force-weight", type=float)
+    parser.add_argument("--stress-weight", type=float)
+    parser.add_argument("--max-train-samples", type=int)
+    parser.add_argument("--max-val-samples", type=int)
+    parser.add_argument("--backend", help="torch.distributed backend for torchrun")
+    parser.add_argument("--seed", type=int)
+    parser.add_argument(
+        "--fit-element-references",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="fit energy element references on the training data",
+    )
+    args = parser.parse_args()
+
+    if not args.config:
+        parser.error("--config is required; use a YAML file from demo/configs")
+    config_path = args.config
+    with Path(config_path).expanduser().open() as handle:
+        config = yaml.safe_load(handle) or {}
+    for key, value in config.items():
+        if getattr(args, key.replace("-", "_"), None) is None:
+            setattr(args, key.replace("-", "_"), value)
+    for key in ("checkpoint", "train", "val", "output"):
+        value = getattr(args, key)
+        if value is not None:
+            if isinstance(value, list):
+                value = [
+                    os.path.expandvars(os.path.expanduser(str(item)))
+                    for item in value
+                ]
+            else:
+                value = os.path.expandvars(os.path.expanduser(str(value)))
+            setattr(args, key, value)
+    required = ("checkpoint", "train", "val", "output")
+    missing = [key for key in required if not getattr(args, key)]
+    if missing:
+        parser.error("missing required config fields: " + ", ".join(missing))
+    args.backend = args.backend or "nccl"
+    args.seed = 0 if args.seed is None else args.seed
+    args.fit_element_references = bool(args.fit_element_references)
+    if not any((args.energy_weight, args.force_weight, args.stress_weight)):
+        parser.error("at least one of energy_weight, force_weight, or stress_weight must be nonzero")
+
+    if args.device.startswith("cuda") and not torch.cuda.is_available():
+        raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false.")
+    context = _init_distributed(args.device, args.backend)
+    try:
+        if args.device.startswith("cuda"):
+            device = torch.device(f"cuda:{context.local_rank}")
+        else:
+            device = torch.device(args.device)
+        torch.manual_seed(args.seed + context.rank)
+
+        # Import registrations before the generic native checkpoint loader.
+        import onescience.models.esen  # noqa: F401
+
+        model = load_model_and_weights_from_checkpoint(args.checkpoint).to(device)
+        transforms = ESENCheckpointTransforms.from_checkpoint(args.checkpoint)
+        if args.fit_element_references:
+            reference_dataset = _loader(
+                args.train, args.batch_size, args.workers
+            ).dataset
+            fitted_references = fit_linear_references(
+                targets=["energy"],
+                dataset=reference_dataset,
+                batch_size=args.batch_size,
+                num_workers=args.workers,
+                log_metrics=False,
+                shuffle=False,
+            )
+            transforms.elementrefs["energy"] = fitted_references["energy"]
+            if context.is_main:
+                print("fitted energy element references from training data", flush=True)
+        transforms = transforms.to(device)
+        if context.enabled:
+            model = DistributedDataParallel(
+                model,
+                device_ids=[context.local_rank] if device.type == "cuda" else None,
+                output_device=context.local_rank if device.type == "cuda" else None,
+            )
+        optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+        train_loader = _loader(
+            args.train,
+            args.batch_size,
+            args.workers,
+            args.max_train_samples,
+            context=context,
+            train=True,
+            seed=args.seed,
+        )
+        val_loader = _loader(
+            args.val,
+            args.batch_size,
+            args.workers,
+            args.max_val_samples,
+            context=context,
+            train=False,
+            seed=args.seed + 1,
+        )
+        weights = (args.energy_weight, args.force_weight, args.stress_weight)
+
+        history = []
+        for epoch in range(args.epochs):
+            if isinstance(train_loader.sampler, DistributedSampler):
+                train_loader.sampler.set_epoch(epoch)
+            train_metrics = _run_epoch(
+                model, train_loader, device, optimizer, weights, transforms, context
+            )
+            # Force/stress outputs are gradients of the energy, so validation also
+            # needs autograd even though model parameters are not updated.
+            val_metrics = _run_epoch(
+                model, val_loader, device, None, weights, transforms, context
+            )
+            record = {"epoch": epoch, "train": train_metrics, "val": val_metrics}
+            if context.is_main:
+                history.append(record)
+                print(json.dumps(record, sort_keys=True), flush=True)
+
+        if context.is_main:
+            output = Path(args.output)
+            output.parent.mkdir(parents=True, exist_ok=True)
+            source = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+            checkpoint = copy.deepcopy(source)
+            base_model = model.module if context.enabled else model
+            checkpoint["state_dict"] = {
+                key: value.detach().cpu() for key, value in base_model.state_dict().items()
+            }
+            checkpoint["elementrefs"] = {
+                name: {
+                    key: value.detach().cpu()
+                    for key, value in elementref.state_dict().items()
+                }
+                for name, elementref in transforms.elementrefs.items()
+            }
+            checkpoint.setdefault("metadata", {})
+            checkpoint["metadata"].update(
+                {
+                    "onescience_esen_history": history,
+                    "source_checkpoint": args.checkpoint,
+                    "world_size": context.world_size,
+                    "loss_space": "checkpoint_normalized",
+                    "element_references": (
+                        "fitted_from_training_data"
+                        if args.fit_element_references
+                        else "source_checkpoint"
+                    ),
+                }
+            )
+            torch.save(checkpoint, output)
+            print(f"saved checkpoint: {output}", flush=True)
+    finally:
+        _close_distributed(context)
+
+
+if __name__ == "__main__":
+    main()

md.py

@@ -0,0 +1,81 @@
+"""Run a short NVT molecular-dynamics trajectory with eSEN."""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+os.environ.setdefault(
+    "ONESCIENCE_ESEN_JD_PATH",
+    os.path.join(os.path.dirname(__file__), "weight", "Jd.pt"),
+)
+
+import numpy as np
+from ase import units
+from ase.build import bulk
+from ase.io import read
+from ase.io.trajectory import Trajectory
+from ase.md.langevin import Langevin
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary
+
+from onescience.utils.esen import eSENCalculator
+
+
+def default_checkpoint() -> str:
+    return os.path.join(os.path.dirname(__file__), "weight", "esen_30m_mptrj.pt")
+
+
+def load_structure(path: str | None, repeat: tuple[int, int, int] | None):
+    if path:
+        atoms = read(path)
+    else:
+        atoms = bulk("Si", "diamond", a=5.43).repeat((2, 2, 2))
+    return atoms.repeat(repeat) if repeat else atoms
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", default=default_checkpoint())
+    parser.add_argument("--input", help="CIF, POSCAR, XYZ, or another ASE-readable structure")
+    parser.add_argument(
+        "--repeat",
+        type=int,
+        nargs=3,
+        metavar=("NX", "NY", "NZ"),
+        help="repeat the input structure along its three cell vectors",
+    )
+    parser.add_argument("--steps", type=int, default=10)
+    parser.add_argument("--temperature", type=float, default=300.0)
+    parser.add_argument("--timestep", type=float, default=1.0, help="time step in fs")
+    parser.add_argument(
+        "--friction", type=float, default=0.01, help="Langevin friction in 1/fs"
+    )
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--output", default="md.traj")
+    parser.add_argument("--device", default="cuda")
+    args = parser.parse_args()
+
+    repeat = tuple(args.repeat) if args.repeat else None
+    atoms = load_structure(args.input, repeat)
+    atoms.calc = eSENCalculator.from_checkpoint(args.checkpoint, device=args.device)
+    rng = np.random.default_rng(args.seed)
+    MaxwellBoltzmannDistribution(atoms, temperature_K=args.temperature, rng=rng)
+    Stationary(atoms)
+    dynamics = Langevin(
+        atoms,
+        timestep=args.timestep * units.fs,
+        temperature_K=args.temperature,
+        friction=args.friction / units.fs,
+    )
+    trajectory = Trajectory(args.output, "w", atoms)
+    dynamics.attach(trajectory.write, interval=1)
+    dynamics.run(args.steps)
+    trajectory.close()
+    print("formula:", atoms.get_chemical_formula())
+    print("atoms:", len(atoms))
+    print("steps:", dynamics.nsteps)
+    print("energy (eV):", atoms.get_potential_energy())
+
+
+if __name__ == "__main__":
+    main()

model/__init__.py

@@ -0,0 +1,18 @@
+"""eSEN atomistic potentials adapted to OneScience's UMA runtime."""
+
+from .esen import eSEN_Backbone, Linear_Force_Head, MLP_EFS_Head, MLP_Energy_Head
+from .esen_dens import (
+    Linear_Force_Head_DeNS,
+    MLP_EFS_Head as MLP_EFS_Head_DeNS,
+    eSEN_DeNS_Backbone,
+)
+
+__all__ = [
+    "eSEN_Backbone",
+    "eSEN_DeNS_Backbone",
+    "MLP_EFS_Head",
+    "MLP_EFS_Head_DeNS",
+    "MLP_Energy_Head",
+    "Linear_Force_Head",
+    "Linear_Force_Head_DeNS",
+]

model/esen.py

@@ -0,0 +1,514 @@
+"""
+Copyright (c) Meta, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn as nn
+
+from onescience.models.UMA.base import BackboneInterface, HeadInterface
+from onescience.models.esen.graph import GraphModelMixin
+from onescience.utils.uma.common.registry import registry
+from onescience.utils.uma.common.utils import conditional_grad
+
+from onescience.modules.layer.esen.common.rotation import (
+    init_edge_rot_mat,
+    rotation_to_wigner,
+)
+from onescience.modules.layer.esen.common.so3 import (
+    CoefficientMapping,
+    SO3_Grid,
+)
+from onescience.modules.layer.esen.esen_block import eSEN_Block
+from onescience.modules.layer.esen.nn.embedding import EdgeDegreeEmbedding
+from onescience.modules.layer.esen.nn.layer_norm import (
+    EquivariantLayerNormArray,
+    EquivariantLayerNormArraySphericalHarmonics,
+    EquivariantRMSNormArraySphericalHarmonicsV2,
+    get_normalization_layer,
+)
+from onescience.modules.layer.esen.nn.radial import EnvelopedBesselBasis, GaussianSmearing
+from onescience.modules.layer.esen.nn.so3_layers import SO3_Linear
+from onescience.modules.func_utils.esen_path_utils import resolve_esen_jd_path
+
+
+@registry.register_model("esen_backbone")
+class eSEN_Backbone(nn.Module, GraphModelMixin):
+    def __init__(
+        self,
+        max_num_elements: int = 100,
+        sphere_channels: int = 128,
+        lmax: int = 2,
+        mmax: int = 2,
+        grid_resolution: int | None = None,
+        otf_graph: bool = False,
+        max_neighbors: int = 300,
+        use_pbc: bool = True,
+        use_pbc_single: bool = False,
+        cutoff: float = 5.0,
+        edge_channels: int = 128,
+        distance_function: str = "gaussian",
+        num_distance_basis: int = 512,
+        direct_forces: bool = True,
+        regress_forces: bool = True,
+        regress_stress: bool = False,
+        # escnmd specific
+        num_layers: int = 2,
+        hidden_channels: int = 128,
+        norm_type: str = "rms_norm_sh",
+        act_type: str = "s2",
+        mlp_type: str = "grid",
+        use_envelope: bool = False,
+        activation_checkpointing: bool = False,
+        jd_path: str | None = None,
+    ):
+        super().__init__()
+
+        self.max_num_elements = max_num_elements
+        self.lmax = lmax
+        self.mmax = mmax
+        self.sphere_channels = sphere_channels
+        self.grid_resolution = grid_resolution
+
+        self.regress_forces = regress_forces
+        self.direct_forces = direct_forces
+        self.regress_stress = regress_stress
+
+        self.otf_graph = otf_graph
+        self.max_neighbors = max_neighbors
+        self.use_pbc = use_pbc
+        self.use_pbc_single = use_pbc_single
+        self.enforce_max_neighbors_strictly = False
+        self.activation_checkpointing = activation_checkpointing
+
+        self.mlp_type = mlp_type
+        self.use_envelope = use_envelope
+
+        # rotation utils
+        Jd_list = torch.load(resolve_esen_jd_path(jd_path))
+        for l in range(self.lmax + 1):
+            self.register_buffer(f"Jd_{l}", Jd_list[l])
+        self.sph_feature_size = int((self.lmax + 1) ** 2)
+        self.mappingReduced = CoefficientMapping(self.lmax, self.mmax)
+
+        # lmax_lmax for node, lmax_mmax for edge
+        self.SO3_grid = nn.ModuleDict()
+        self.SO3_grid["lmax_lmax"] = SO3_Grid(
+            self.lmax, self.lmax, resolution=grid_resolution, rescale=True
+        )
+        self.SO3_grid["lmax_mmax"] = SO3_Grid(
+            self.lmax, self.mmax, resolution=grid_resolution, rescale=True
+        )
+
+        # atom embedding
+        self.sphere_embedding = nn.Embedding(
+            self.max_num_elements, self.sphere_channels
+        )
+
+        # edge distance embedding
+        self.cutoff = cutoff
+        self.edge_channels = edge_channels
+        self.distance_function = distance_function
+        self.num_distance_basis = num_distance_basis
+
+        if self.distance_function == "gaussian":
+            self.distance_expansion = GaussianSmearing(
+                0.0,
+                self.cutoff,
+                self.num_distance_basis,
+                2.0,
+            )
+        elif self.distance_function == "bessel":
+            self.distance_expansion = EnvelopedBesselBasis(
+                num_radial=self.num_distance_basis,
+                cutoff=cutoff,
+            )
+            self.distance_expansion.offset = [self.cutoff]
+            self.distance_expansion.num_output = self.num_distance_basis
+        else:
+            raise ValueError("Unknown distance function")
+
+        # equivariant initial embedding
+        self.source_embedding = nn.Embedding(self.max_num_elements, self.edge_channels)
+        self.target_embedding = nn.Embedding(self.max_num_elements, self.edge_channels)
+        nn.init.uniform_(self.source_embedding.weight.data, -0.001, 0.001)
+        nn.init.uniform_(self.target_embedding.weight.data, -0.001, 0.001)
+
+        self.edge_channels_list = [
+            self.num_distance_basis + 2 * self.edge_channels,
+            self.edge_channels,
+            self.edge_channels,
+        ]
+
+        self.edge_degree_embedding = EdgeDegreeEmbedding(
+            sphere_channels=self.sphere_channels,
+            lmax=self.lmax,
+            mmax=self.mmax,
+            max_num_elements=self.max_num_elements,
+            edge_channels_list=self.edge_channels_list,
+            rescale_factor=5.0,
+            cutoff=self.cutoff,
+            mappingReduced=self.mappingReduced,
+            out_mask=self.SO3_grid["lmax_lmax"].mapping.coefficient_idx(
+                self.lmax, self.mmax
+            ),
+            use_envelope=use_envelope,
+        )
+
+        self.num_layers = num_layers
+        self.hidden_channels = hidden_channels
+        self.norm_type = norm_type
+        self.act_type = act_type
+
+        # Initialize the blocks for each layer
+        self.blocks = nn.ModuleList()
+        for _ in range(self.num_layers):
+            block = eSEN_Block(
+                self.sphere_channels,
+                self.hidden_channels,
+                self.lmax,
+                self.mmax,
+                self.mappingReduced,
+                self.SO3_grid,
+                self.edge_channels_list,
+                self.cutoff,
+                self.norm_type,
+                self.act_type,
+                self.mlp_type,
+                self.use_envelope,
+            )
+            self.blocks.append(block)
+
+        self.norm = get_normalization_layer(
+            self.norm_type,
+            lmax=self.lmax,
+            num_channels=self.sphere_channels,
+        )
+
+    def get_rotmat_and_wigner(self, edge_distance_vecs):
+        edge_rot_mat = init_edge_rot_mat(
+            edge_distance_vecs, rot_clip=(not self.direct_forces)
+        )
+
+        Jd_buffers = [
+            getattr(self, f"Jd_{l}").type(edge_rot_mat.dtype)
+            for l in range(self.lmax + 1)
+        ]
+
+        wigner = rotation_to_wigner(
+            edge_rot_mat,
+            0,
+            self.lmax,
+            Jd_buffers,
+            rot_clip=(not self.direct_forces),
+        )
+        wigner_inv = torch.transpose(wigner, 1, 2).contiguous()
+
+        return edge_rot_mat, wigner, wigner_inv
+
+    def generate_graph(self, *args, **kwargs):
+        graph = super().generate_graph(*args, **kwargs)
+        return {
+            "edge_index": graph.edge_index,
+            "edge_distance": graph.edge_distance,
+            "edge_distance_vec": graph.edge_distance_vec,
+            "cell_offsets": graph.cell_offsets,
+            "offset_distances": None,
+            "neighbors": None,
+            "node_offset": 0,
+            "batch_full": graph.batch_full,
+            "atomic_numbers_full": graph.atomic_numbers_full,
+        }
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data_dict) -> dict[str, torch.Tensor]:
+        ###############################################################
+        # gradient-based forces/stress
+        ###############################################################
+        data_dict["atomic_numbers"] = data_dict["atomic_numbers"].long()
+
+        displacement = None
+        orig_cell = None
+        if self.regress_stress and not self.direct_forces:
+            displacement = torch.zeros(
+                (3, 3),
+                dtype=data_dict["pos"].dtype,
+                device=data_dict["pos"].device,
+            )
+            # num_batch = data_dict["num_graphs"]
+            num_batch = data_dict.get("num_graphs", len(data_dict["natoms"]))
+            displacement = displacement.view(-1, 3, 3).expand(num_batch, 3, 3)
+            displacement.requires_grad_(True)
+            symmetric_displacement = 0.5 * (
+                displacement + displacement.transpose(-1, -2)
+            )
+
+            data_dict["pos"].requires_grad_(True)
+            data_dict["pos"] = data_dict["pos"] + torch.bmm(
+                data_dict["pos"].unsqueeze(-2),
+                torch.index_select(symmetric_displacement, 0, data_dict["batch"]),
+            ).squeeze(-2)
+
+            orig_cell = data_dict["cell"]
+            data_dict["cell"] = data_dict["cell"] + torch.bmm(
+                data_dict["cell"], symmetric_displacement
+            )
+
+        if not self.regress_stress and self.regress_forces and not self.direct_forces:
+            data_dict["pos"].requires_grad_(True)
+
+        if self.otf_graph:
+            graph_dict = self.generate_graph(data_dict)
+        else:
+            cell_per_edge = data_dict["cell"].repeat_interleave(
+                data_dict["nedges"], dim=0
+            )
+            shifts = torch.einsum(
+                "ij,ijk->ik",
+                data_dict["cell_offsets"].to(cell_per_edge.dtype),
+                cell_per_edge,
+            )
+            edge_distance_vec = (
+                data_dict["pos"][data_dict["edge_index"][0]]
+                - data_dict["pos"][data_dict["edge_index"][1]]
+                + shifts
+            )
+            # pylint: disable=E1102
+            edge_distance = torch.linalg.norm(edge_distance_vec, dim=-1, keepdim=False)
+            graph_dict = {
+                "atomic_numbers_full": data_dict["atomic_numbers_full"],
+                "batch_full": data_dict["batch_full"],
+                "edge_index": data_dict["edge_index"],
+                "edge_distance": edge_distance,
+                "edge_distance_vec": edge_distance_vec,
+                "node_offset": 0,
+            }
+
+        _, wigner, wigner_inv = self.get_rotmat_and_wigner(
+            graph_dict["edge_distance_vec"]
+        )
+
+        ###############################################################
+        # Initialize node embeddings
+        ###############################################################
+
+        x_message = torch.zeros(
+            data_dict["pos"].shape[0],
+            self.sph_feature_size,
+            self.sphere_channels,
+            device=data_dict["pos"].device,
+            dtype=data_dict["pos"].dtype,
+        )
+        x_message[:, 0, :] = self.sphere_embedding(data_dict["atomic_numbers"])
+
+        # edge degree embedding
+        edge_distance_embedding = self.distance_expansion(graph_dict["edge_distance"])
+        source_embedding = self.source_embedding(
+            data_dict["atomic_numbers"][graph_dict["edge_index"][0]]
+        )
+        target_embedding = self.target_embedding(
+            data_dict["atomic_numbers"][graph_dict["edge_index"][1]]
+        )
+        x_edge = torch.cat(
+            (edge_distance_embedding, source_embedding, target_embedding), dim=1
+        )
+        x_message = self.edge_degree_embedding(
+            x_message,
+            x_edge,
+            graph_dict["edge_distance"],
+            graph_dict["edge_index"],
+            wigner_inv,
+        )
+
+        ###############################################################
+        # Update spherical node embeddings
+        ###############################################################
+        if graph_dict["edge_index"].shape[1] != 0:
+            for i in range(self.num_layers):
+                if self.activation_checkpointing:
+                    x_message = torch.utils.checkpoint.checkpoint(
+                        self.blocks[i],
+                        x_message,
+                        x_edge,
+                        graph_dict["edge_distance"],
+                        graph_dict["edge_index"],
+                        wigner,
+                        wigner_inv,
+                        graph_dict["node_offset"],
+                        use_reentrant=False,
+                    )
+                else:
+                    x_message = self.blocks[i](
+                        x_message,
+                        x_edge,
+                        graph_dict["edge_distance"],
+                        graph_dict["edge_index"],
+                        wigner,
+                        wigner_inv,
+                        node_offset=graph_dict["node_offset"],
+                    )
+
+        # Final layer norm
+        x_message = self.norm(x_message)
+
+        out = {
+            "node_embedding": x_message,
+            "displacement": displacement,
+            "orig_cell": orig_cell,
+        }
+        out.update(graph_dict)
+
+        return out
+
+    @property
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    @torch.jit.ignore
+    def no_weight_decay(self) -> set:
+        no_wd_list = []
+        named_parameters_list = [name for name, _ in self.named_parameters()]
+        for module_name, module in self.named_modules():
+            if isinstance(
+                module,
+                (
+                    torch.nn.Linear,
+                    SO3_Linear,
+                    torch.nn.LayerNorm,
+                    EquivariantLayerNormArray,
+                    EquivariantLayerNormArraySphericalHarmonics,
+                    EquivariantRMSNormArraySphericalHarmonicsV2,
+                ),
+            ):
+                for parameter_name, _ in module.named_parameters():
+                    if (
+                        isinstance(module, (torch.nn.Linear, SO3_Linear))
+                        and "weight" in parameter_name
+                    ):
+                        continue
+                    global_parameter_name = module_name + "." + parameter_name
+                    assert global_parameter_name in named_parameters_list
+                    no_wd_list.append(global_parameter_name)
+
+        return set(no_wd_list)
+
+
+@registry.register_model("esen_mlp_efs_head")
+class MLP_EFS_Head(nn.Module, HeadInterface):
+    def __init__(self, backbone):
+        super().__init__()
+        backbone.energy_block = None
+        backbone.force_block = None
+        self.regress_stress = backbone.regress_stress
+        self.regress_forces = backbone.regress_forces
+
+        self.sphere_channels = backbone.sphere_channels
+        self.hidden_channels = backbone.hidden_channels
+        self.energy_block = nn.Sequential(
+            nn.Linear(self.sphere_channels, self.hidden_channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(self.hidden_channels, self.hidden_channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(self.hidden_channels, 1, bias=True),
+        )
+
+        backbone.direct_forces = False
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data, emb: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        energy_key = "energy"
+        forces_key = "forces"
+        stress_key = "stress"
+
+        outputs = {}
+
+        node_energy = self.energy_block(
+            emb["node_embedding"].narrow(1, 0, 1).squeeze()
+        ).view(-1, 1, 1)
+
+        energy = torch.zeros(
+            len(data["natoms"]), device=data["pos"].device, dtype=node_energy.dtype
+        )
+        energy.index_add_(0, data["batch"], node_energy.view(-1))
+        outputs[energy_key] = energy
+
+        if self.regress_stress:
+            grads = torch.autograd.grad(
+                [energy.sum()],
+                [data["pos"], emb["displacement"]],
+                create_graph=self.training,
+            )
+            forces = torch.neg(grads[0])
+            virial = grads[1].view(-1, 3, 3)
+            volume = torch.det(data["cell"]).abs().unsqueeze(-1)
+            stress = virial / volume.view(-1, 1, 1)
+            virial = torch.neg(virial)
+            outputs[forces_key] = forces
+            outputs[stress_key] = stress.view(-1, 9)
+            data["cell"] = emb["orig_cell"]
+        elif self.regress_forces:
+            forces = (
+                -1
+                * torch.autograd.grad(
+                    energy.sum(), data["pos"], create_graph=self.training
+                )[0]
+            )
+            outputs[forces_key] = forces
+        return outputs
+
+
+@registry.register_model("esen_mlp_energy_head")
+class MLP_Energy_Head(nn.Module, HeadInterface):
+    def __init__(self, backbone, reduce: str = "sum"):
+        super().__init__()
+        self.reduce = reduce
+
+        self.sphere_channels = backbone.sphere_channels
+        self.hidden_channels = backbone.hidden_channels
+        self.energy_block = nn.Sequential(
+            nn.Linear(self.sphere_channels, self.hidden_channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(self.hidden_channels, self.hidden_channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(self.hidden_channels, 1, bias=True),
+        )
+
+    def forward(self, data_dict, emb: dict[str, torch.Tensor]):
+        node_energy = self.energy_block(
+            emb["node_embedding"].narrow(1, 0, 1).squeeze()
+        ).view(-1, 1, 1)
+
+        energy = torch.zeros(
+            len(data_dict["natoms"]),
+            device=node_energy.device,
+            dtype=node_energy.dtype,
+        )
+
+        energy.index_add_(0, data_dict["batch"], node_energy.view(-1))
+        if self.reduce == "sum":
+            return {"energy": energy}
+        elif self.reduce == "mean":
+            return {"energy": energy / data_dict["natoms"]}
+        else:
+            raise ValueError(
+                f"reduce can only be sum or mean, user provided: {self.reduce}"
+            )
+
+
+@registry.register_model("esen_linear_force_head")
+class Linear_Force_Head(nn.Module, HeadInterface):
+    def __init__(self, backbone):
+        super().__init__()
+        self.linear = SO3_Linear(backbone.sphere_channels, 1, lmax=1)
+
+    def forward(self, data_dict, emb: dict[str, torch.Tensor]):
+        forces = self.linear(emb["node_embedding"].narrow(1, 0, 4))
+        forces = forces.narrow(1, 1, 3)
+        forces = forces.view(-1, 3).contiguous()
+        return {"forces": forces}

model/esen_dens.py

@@ -0,0 +1,582 @@
+"""
+Copyright (c) Meta, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import torch.nn as nn
+from e3nn import o3
+
+from onescience.models.UMA.base import BackboneInterface, HeadInterface
+from onescience.models.esen.graph import GraphModelMixin
+from onescience.utils.uma.common.registry import registry
+from onescience.utils.uma.common.utils import conditional_grad
+
+from onescience.modules.layer.esen.common.rotation import (
+    init_edge_rot_mat,
+    rotation_to_wigner,
+)
+from onescience.modules.layer.esen.common.so3 import (
+    CoefficientMapping,
+    SO3_Grid,
+)
+from onescience.modules.layer.esen.esen_block import eSEN_Block
+from onescience.modules.layer.esen.nn.embedding import EdgeDegreeEmbedding
+from onescience.modules.layer.esen.nn.layer_norm import (
+    EquivariantLayerNormArray,
+    EquivariantLayerNormArraySphericalHarmonics,
+    EquivariantRMSNormArraySphericalHarmonicsV2,
+    get_normalization_layer,
+)
+from onescience.modules.layer.esen.nn.radial import EnvelopedBesselBasis, GaussianSmearing
+from onescience.modules.layer.esen.nn.so3_layers import SO3_Linear
+from onescience.modules.func_utils.esen_path_utils import resolve_esen_jd_path
+
+
+@registry.register_model("esen_backbone_dens")
+class eSEN_DeNS_Backbone(nn.Module, GraphModelMixin):
+    def __init__(
+        self,
+        max_num_elements: int = 100,
+        sphere_channels: int = 128,
+        lmax: int = 2,
+        mmax: int = 2,
+        grid_resolution: int | None = None,
+        otf_graph: bool = False,
+        max_neighbors: int = 300,
+        use_pbc: bool = True,
+        use_pbc_single: bool = False,
+        cutoff: float = 5.0,
+        edge_channels: int = 128,
+        distance_function: str = "gaussian",
+        num_distance_basis: int = 512,
+        direct_forces: bool = True,
+        regress_forces: bool = True,
+        regress_stress: bool = False,
+        # escnmd specific
+        num_layers: int = 2,
+        hidden_channels: int = 128,
+        norm_type: str = "rms_norm_sh",
+        act_type: str = "s2",
+        mlp_type: str = "grid",
+        use_envelope: bool = False,
+        activation_checkpointing: bool = False,
+        jd_path: str | None = None,
+    ):
+        super().__init__()
+
+        self.max_num_elements = max_num_elements
+        self.lmax = lmax
+        self.mmax = mmax
+        self.sphere_channels = sphere_channels
+        self.grid_resolution = grid_resolution
+
+        self.regress_forces = regress_forces
+        self.direct_forces = direct_forces
+        self.regress_stress = regress_stress
+
+        self.otf_graph = otf_graph
+        self.max_neighbors = max_neighbors
+        self.use_pbc = use_pbc
+        self.use_pbc_single = use_pbc_single
+        self.enforce_max_neighbors_strictly = False
+        self.activation_checkpointing = activation_checkpointing
+
+        self.mlp_type = mlp_type
+        self.use_envelope = use_envelope
+
+        # rotation utils
+        Jd_list = torch.load(resolve_esen_jd_path(jd_path))
+        for l in range(self.lmax + 1):
+            self.register_buffer(f"Jd_{l}", Jd_list[l])
+        self.sph_feature_size = int((self.lmax + 1) ** 2)
+        self.mappingReduced = CoefficientMapping(self.lmax, self.mmax)
+
+        # lmax_lmax for node, lmax_mmax for edge
+        self.SO3_grid = nn.ModuleDict()
+        self.SO3_grid["lmax_lmax"] = SO3_Grid(
+            self.lmax, self.lmax, resolution=grid_resolution, rescale=True
+        )
+        self.SO3_grid["lmax_mmax"] = SO3_Grid(
+            self.lmax, self.mmax, resolution=grid_resolution, rescale=True
+        )
+
+        # atom embedding
+        self.sphere_embedding = nn.Embedding(
+            self.max_num_elements, self.sphere_channels
+        )
+
+        # edge distance embedding
+        self.cutoff = cutoff
+        self.edge_channels = edge_channels
+        self.distance_function = distance_function
+        self.num_distance_basis = num_distance_basis
+
+        if self.distance_function == "gaussian":
+            self.distance_expansion = GaussianSmearing(
+                0.0,
+                self.cutoff,
+                self.num_distance_basis,
+                2.0,
+            )
+        elif self.distance_function == "bessel":
+            self.distance_expansion = EnvelopedBesselBasis(
+                num_radial=self.num_distance_basis,
+                cutoff=cutoff,
+            )
+            self.distance_expansion.offset = [self.cutoff]
+            self.distance_expansion.num_output = self.num_distance_basis
+        else:
+            raise ValueError("Unknown distance function")
+
+        # equivariant initial embedding
+        self.source_embedding = nn.Embedding(self.max_num_elements, self.edge_channels)
+        self.target_embedding = nn.Embedding(self.max_num_elements, self.edge_channels)
+        nn.init.uniform_(self.source_embedding.weight.data, -0.001, 0.001)
+        nn.init.uniform_(self.target_embedding.weight.data, -0.001, 0.001)
+
+        self.edge_channels_list = [
+            self.num_distance_basis + 2 * self.edge_channels,
+            self.edge_channels,
+            self.edge_channels,
+        ]
+
+        self.edge_degree_embedding = EdgeDegreeEmbedding(
+            sphere_channels=self.sphere_channels,
+            lmax=self.lmax,
+            mmax=self.mmax,
+            max_num_elements=self.max_num_elements,
+            edge_channels_list=self.edge_channels_list,
+            rescale_factor=5.0,
+            cutoff=self.cutoff,
+            mappingReduced=self.mappingReduced,
+            out_mask=self.SO3_grid["lmax_lmax"].mapping.coefficient_idx(
+                self.lmax, self.mmax
+            ),
+            use_envelope=use_envelope,
+        )
+
+        self.num_layers = num_layers
+        self.hidden_channels = hidden_channels
+        self.norm_type = norm_type
+        self.act_type = act_type
+
+        # Initialize the blocks for each layer
+        self.blocks = nn.ModuleList()
+        for _ in range(self.num_layers):
+            block = eSEN_Block(
+                self.sphere_channels,
+                self.hidden_channels,
+                self.lmax,
+                self.mmax,
+                self.mappingReduced,
+                self.SO3_grid,
+                self.edge_channels_list,
+                self.cutoff,
+                self.norm_type,
+                self.act_type,
+                self.mlp_type,
+                self.use_envelope,
+            )
+            self.blocks.append(block)
+
+        self.norm = get_normalization_layer(
+            self.norm_type,
+            lmax=self.lmax,
+            num_channels=self.sphere_channels,
+        )
+
+        self.irreps_sh = o3.Irreps.spherical_harmonics(lmax=self.lmax, p=1)
+        self.force_embedding = SO3_Linear(
+            in_features=1, out_features=self.sphere_channels, lmax=self.lmax
+        )
+
+    def get_rotmat_and_wigner(self, edge_distance_vecs):
+        edge_rot_mat = init_edge_rot_mat(
+            edge_distance_vecs, rot_clip=(not self.direct_forces)
+        )
+
+        Jd_buffers = [
+            getattr(self, f"Jd_{l}").type(edge_rot_mat.dtype)
+            for l in range(self.lmax + 1)
+        ]
+
+        wigner = rotation_to_wigner(
+            edge_rot_mat,
+            0,
+            self.lmax,
+            Jd_buffers,
+            rot_clip=(not self.direct_forces),
+        )
+        wigner_inv = torch.transpose(wigner, 1, 2).contiguous()
+
+        return edge_rot_mat, wigner, wigner_inv
+
+    def generate_graph(self, *args, **kwargs):
+        graph = super().generate_graph(*args, **kwargs)
+        return {
+            "edge_index": graph.edge_index,
+            "edge_distance": graph.edge_distance,
+            "edge_distance_vec": graph.edge_distance_vec,
+            "cell_offsets": graph.cell_offsets,
+            "offset_distances": None,
+            "neighbors": None,
+            "node_offset": 0,
+            "batch_full": graph.batch_full,
+            "atomic_numbers_full": graph.atomic_numbers_full,
+        }
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data_dict) -> dict[str, torch.Tensor]:
+        ###############################################################
+        # gradient-based forces/stress
+        ###############################################################
+        data_dict["atomic_numbers"] = data_dict["atomic_numbers"].long()
+
+        displacement = None
+        orig_cell = None
+        if self.regress_stress and not self.direct_forces:
+            displacement = torch.zeros(
+                (3, 3),
+                dtype=data_dict["pos"].dtype,
+                device=data_dict["pos"].device,
+            )
+            # num_batch = data_dict["num_graphs"]
+            num_batch = data_dict.get("num_graphs", len(data_dict["natoms"]))
+            displacement = displacement.view(-1, 3, 3).expand(num_batch, 3, 3)
+            displacement.requires_grad_(True)
+            symmetric_displacement = 0.5 * (
+                displacement + displacement.transpose(-1, -2)
+            )
+
+            data_dict["pos"].requires_grad_(True)
+            data_dict["pos"] = data_dict["pos"] + torch.bmm(
+                data_dict["pos"].unsqueeze(-2),
+                torch.index_select(symmetric_displacement, 0, data_dict["batch"]),
+            ).squeeze(-2)
+
+            orig_cell = data_dict["cell"]
+            data_dict["cell"] = data_dict["cell"] + torch.bmm(
+                data_dict["cell"], symmetric_displacement
+            )
+
+        if not self.regress_stress and self.regress_forces and not self.direct_forces:
+            data_dict["pos"].requires_grad_(True)
+
+        if self.otf_graph:
+            graph_dict = self.generate_graph(data_dict)
+        else:
+            cell_per_edge = data_dict["cell"].repeat_interleave(
+                data_dict["nedges"], dim=0
+            )
+            shifts = torch.einsum(
+                "ij,ijk->ik",
+                data_dict["cell_offsets"].to(cell_per_edge.dtype),
+                cell_per_edge,
+            )
+            edge_distance_vec = (
+                data_dict["pos"][data_dict["edge_index"][0]]
+                - data_dict["pos"][data_dict["edge_index"][1]]
+                + shifts
+            )
+            # pylint: disable=E1102
+            edge_distance = torch.linalg.norm(edge_distance_vec, dim=-1, keepdim=False)
+            graph_dict = {
+                "atomic_numbers_full": data_dict["atomic_numbers_full"],
+                "batch_full": data_dict["batch_full"],
+                "edge_index": data_dict["edge_index"],
+                "edge_distance": edge_distance,
+                "edge_distance_vec": edge_distance_vec,
+                "node_offset": 0,
+            }
+
+        _, wigner, wigner_inv = self.get_rotmat_and_wigner(
+            graph_dict["edge_distance_vec"]
+        )
+
+        ###############################################################
+        # Initialize node embeddings
+        ###############################################################
+
+        x_message = torch.zeros(
+            data_dict["pos"].shape[0],
+            self.sph_feature_size,
+            self.sphere_channels,
+            device=data_dict["pos"].device,
+            dtype=data_dict["pos"].dtype,
+        )
+        x_message[:, 0, :] = self.sphere_embedding(data_dict["atomic_numbers"])
+
+        ##################
+        ### DeNS Start ###
+        ##################
+
+        num_atoms = len(data_dict.atomic_numbers)
+        # Node-wise force encoding during denoising positions
+        if (
+            hasattr(data_dict, "denoising_pos_forward")
+            and data_dict.denoising_pos_forward
+        ):
+            assert hasattr(data_dict, "forces")
+            force_data = data_dict.forces
+            force_sh = o3.spherical_harmonics(
+                l=self.irreps_sh,
+                x=force_data,
+                normalize=True,
+                normalization="component",
+            )
+            force_sh = force_sh.view(num_atoms, (self.lmax + 1) ** 2, 1)
+            force_norm = force_data.norm(dim=-1, keepdim=True)
+            if hasattr(data_dict, "noise_mask"):
+                noise_mask_tensor = data_dict.noise_mask.view(-1, 1, 1)
+                force_sh = force_sh * noise_mask_tensor
+        else:
+            force_sh = torch.zeros(
+                (num_atoms, (self.lmax + 1) ** 2, 1),
+                dtype=data_dict.pos.dtype,
+                device=data_dict.pos.device,
+            )
+            force_norm = torch.zeros(
+                (num_atoms, 1), dtype=data_dict.pos.dtype, device=data_dict.pos.device
+            )
+
+        force_norm = force_norm.view(-1, 1, 1)
+        force_norm = force_norm / math.sqrt(
+            3.0
+        )  # since we use `component` normalization
+        force_embedding = force_sh * force_norm
+
+        force_embedding = self.force_embedding(force_embedding)
+        x_message = x_message + force_embedding
+
+        ##################
+        ###  DeNS End  ###
+        ##################
+
+        # edge degree embedding
+        edge_distance_embedding = self.distance_expansion(graph_dict["edge_distance"])
+        source_embedding = self.source_embedding(
+            data_dict["atomic_numbers"][graph_dict["edge_index"][0]]
+        )
+        target_embedding = self.target_embedding(
+            data_dict["atomic_numbers"][graph_dict["edge_index"][1]]
+        )
+        x_edge = torch.cat(
+            (edge_distance_embedding, source_embedding, target_embedding), dim=1
+        )
+        x_message = self.edge_degree_embedding(
+            x_message,
+            x_edge,
+            graph_dict["edge_distance"],
+            graph_dict["edge_index"],
+            wigner_inv,
+        )
+
+        ###############################################################
+        # Update spherical node embeddings
+        ###############################################################
+        if graph_dict["edge_index"].shape[1] != 0:
+            for i in range(self.num_layers):
+                if self.activation_checkpointing:
+                    x_message = torch.utils.checkpoint.checkpoint(
+                        self.blocks[i],
+                        x_message,
+                        x_edge,
+                        graph_dict["edge_distance"],
+                        graph_dict["edge_index"],
+                        wigner,
+                        wigner_inv,
+                        graph_dict["node_offset"],
+                        use_reentrant=False,
+                    )
+                else:
+                    x_message = self.blocks[i](
+                        x_message,
+                        x_edge,
+                        graph_dict["edge_distance"],
+                        graph_dict["edge_index"],
+                        wigner,
+                        wigner_inv,
+                        node_offset=graph_dict["node_offset"],
+                    )
+
+        # Final layer norm
+        x_message = self.norm(x_message)
+
+        out = {
+            "node_embedding": x_message,
+            "displacement": displacement,
+            "orig_cell": orig_cell,
+        }
+        out.update(graph_dict)
+
+        return out
+
+    @property
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    @torch.jit.ignore
+    def no_weight_decay(self) -> set:
+        no_wd_list = []
+        named_parameters_list = [name for name, _ in self.named_parameters()]
+        for module_name, module in self.named_modules():
+            if isinstance(
+                module,
+                (
+                    torch.nn.Linear,
+                    SO3_Linear,
+                    torch.nn.LayerNorm,
+                    EquivariantLayerNormArray,
+                    EquivariantLayerNormArraySphericalHarmonics,
+                    EquivariantRMSNormArraySphericalHarmonicsV2,
+                ),
+            ):
+                for parameter_name, _ in module.named_parameters():
+                    if (
+                        isinstance(module, (torch.nn.Linear, SO3_Linear))
+                        and "weight" in parameter_name
+                    ):
+                        continue
+                    global_parameter_name = module_name + "." + parameter_name
+                    assert global_parameter_name in named_parameters_list
+                    no_wd_list.append(global_parameter_name)
+
+        return set(no_wd_list)
+
+
+@registry.register_model("esen_mlp_efs_head_dens")
+class MLP_EFS_Head(nn.Module, HeadInterface):
+    def __init__(self, backbone):
+        super().__init__()
+        backbone.energy_block = None
+        backbone.force_block = None
+        self.regress_stress = backbone.regress_stress
+        self.regress_forces = backbone.regress_forces
+        self.denoising_pos = True
+
+        self.sphere_channels = backbone.sphere_channels
+        self.hidden_channels = backbone.hidden_channels
+        self.energy_block = nn.Sequential(
+            nn.Linear(self.sphere_channels, self.hidden_channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(self.hidden_channels, self.hidden_channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(self.hidden_channels, 1, bias=True),
+        )
+
+        self.denoising_linear = SO3_Linear(backbone.sphere_channels, 1, lmax=1)
+
+        backbone.direct_forces = False
+
+    def delete_denoising_head(self):
+        self.denoising_linear = None
+        self.denoising_pos = False
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data, emb: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        energy_key = "energy"
+        forces_key = "forces"
+        stress_key = "stress"
+
+        outputs = {}
+
+        node_energy = self.energy_block(
+            emb["node_embedding"].narrow(1, 0, 1).squeeze()
+        ).view(-1, 1, 1)
+
+        energy = torch.zeros(
+            len(data["natoms"]), device=data["pos"].device, dtype=node_energy.dtype
+        )
+        energy.index_add_(0, data["batch"], node_energy.view(-1))
+
+        # no train on noisy energy for conserved
+        if hasattr(data, "denoising_pos_forward") and data.denoising_pos_forward:
+            energy = energy * 0.0
+
+        outputs[energy_key] = energy
+
+        if self.regress_stress:
+            grads = torch.autograd.grad(
+                [energy.sum()],
+                [data["pos"], emb["displacement"]],
+                create_graph=self.training,
+            )
+            forces = torch.neg(grads[0])
+            virial = grads[1].view(-1, 3, 3)
+            volume = torch.det(data["cell"]).abs().unsqueeze(-1)
+            stress = virial / volume.view(-1, 1, 1)
+            virial = torch.neg(virial)
+            outputs[forces_key] = forces
+            outputs[stress_key] = stress.view(
+                -1, 9
+            )  # NOTE to work better with current Multi-task trainer
+            data["cell"] = emb["orig_cell"]
+        elif self.regress_forces:
+            forces = (
+                -1
+                * torch.autograd.grad(
+                    energy.sum(), data["pos"], create_graph=self.training
+                )[0]
+            )
+
+        if self.denoising_pos:
+            denoising_pos_vec = self.denoising_linear(
+                emb["node_embedding"].narrow(1, 0, 4)
+            )
+            denoising_pos_vec = denoising_pos_vec.narrow(1, 1, 3)
+            denoising_pos_vec = denoising_pos_vec.view(-1, 3).contiguous()
+
+            if hasattr(data, "denoising_pos_forward") and data.denoising_pos_forward:
+                if hasattr(data, "noise_mask"):
+                    noise_mask_tensor = data.noise_mask.view(-1, 1)
+                    forces = denoising_pos_vec * noise_mask_tensor + forces * (
+                        ~noise_mask_tensor
+                    )
+                else:
+                    forces = denoising_pos_vec + 0 * forces
+            else:
+                forces = 0 * denoising_pos_vec + forces
+
+        outputs[forces_key] = forces
+
+        return outputs
+
+
+@registry.register_model("esen_linear_force_head_dens")
+class Linear_Force_Head_DeNS(nn.Module, HeadInterface):
+    def __init__(self, backbone):
+        super().__init__()
+        self.linear = SO3_Linear(backbone.sphere_channels, 1, lmax=1)
+        self.denoising_linear = SO3_Linear(backbone.sphere_channels, 1, lmax=1)
+
+    def forward(self, data_dict, emb: dict[str, torch.Tensor]):
+        forces = self.linear(emb["node_embedding"].narrow(1, 0, 4))
+        denoising_pos_vec = self.denoising_linear(emb["node_embedding"].narrow(1, 0, 4))
+        forces = forces.narrow(1, 1, 3)
+        forces = forces.view(-1, 3).contiguous()
+        denoising_pos_vec = denoising_pos_vec.narrow(1, 1, 3)
+        denoising_pos_vec = denoising_pos_vec.view(-1, 3).contiguous()
+
+        if (
+            hasattr(data_dict, "denoising_pos_forward")
+            and data_dict.denoising_pos_forward
+        ):
+            if hasattr(data_dict, "noise_mask"):
+                noise_mask_tensor = data_dict.noise_mask.view(-1, 1)
+                forces = denoising_pos_vec * noise_mask_tensor + forces * (
+                    ~noise_mask_tensor
+                )
+            else:
+                forces = denoising_pos_vec + 0 * forces
+        else:
+            forces = 0 * denoising_pos_vec + forces
+
+        return {"forces": forces}

model/graph.py

@@ -0,0 +1,43 @@
+"""Graph helpers shared by the OneScience eSEN backbones."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+
+from onescience.modules.func_utils.uma_graph.compute import generate_graph
+
+
+class GraphModelMixin:
+    """Provide FairChem-v1 compatible graph output using OneScience graph code."""
+
+    def generate_graph(self, data, cutoff=None, max_neighbors=None, **kwargs):
+        cutoff = cutoff or self.cutoff
+        max_neighbors = max_neighbors or self.max_neighbors
+        pbc = kwargs.pop("pbc", None)
+        if pbc is None:
+            pbc = getattr(data, "pbc", None)
+        if pbc is None:
+            pbc = torch.ones(
+                (data.natoms.numel(), 3), dtype=torch.bool, device=data.pos.device
+            )
+        elif pbc.ndim == 1:
+            pbc = pbc.view(1, 3).expand(data.natoms.numel(), -1)
+
+        graph = generate_graph(
+            data,
+            cutoff=cutoff,
+            max_neighbors=max_neighbors,
+            enforce_max_neighbors_strictly=getattr(
+                self, "enforce_max_neighbors_strictly", False
+            ),
+            radius_pbc_version=getattr(self, "radius_pbc_version", 1),
+            pbc=pbc,
+        )
+        graph.update(
+            batch_full=data.batch,
+            atomic_numbers_full=data.atomic_numbers,
+            node_offset=0,
+        )
+        return SimpleNamespace(**graph)

prepare_oxide_dataset.py

@@ -0,0 +1,114 @@
+"""Build reproducible ASE DB splits from FairChem's oxide tutorial data."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+from ase import Atoms, units
+from ase.calculators.singlepoint import SinglePointCalculator
+from ase.db import connect
+
+
+def load_records(path: Path):
+    with path.open() as handle:
+        source = json.load(handle)
+    records = []
+    for oxide, polymorphs in source.items():
+        for polymorph, calculations in polymorphs.items():
+            for calculation in calculations.get("PBE", {}).get("EOS", {}).get(
+                "calculations", []
+            ):
+                records.append((oxide, polymorph, calculation))
+    return records
+
+
+def split_groups(records, seed: int):
+    groups = sorted({(oxide, polymorph) for oxide, polymorph, _ in records})
+    random.Random(seed).shuffle(groups)
+    n_train = int(0.8 * len(groups))
+    n_val = int(0.1 * len(groups))
+    split_for = {group: "train" for group in groups[:n_train]}
+    split_for.update({group: "val" for group in groups[n_train : n_train + n_val]})
+    split_for.update({group: "test" for group in groups[n_train + n_val :]})
+    return split_for
+
+
+def to_atoms(oxide: str, polymorph: str, calculation: dict) -> Atoms:
+    structure = calculation["atoms"]
+    results = calculation["data"]
+    atoms = Atoms(
+        symbols=structure["symbols"],
+        positions=structure["positions"],
+        cell=structure["cell"],
+        pbc=structure["pbc"],
+    )
+    atoms.set_tags(np.ones(len(atoms), dtype=int))
+    # ASE stress uses eV/Angstrom^3. The source JSON documents stress in GPa.
+    stress = np.asarray(results["stress"], dtype=float) * units.GPa
+    atoms.calc = SinglePointCalculator(
+        atoms,
+        energy=float(results["total_energy"]),
+        forces=np.asarray(results["forces"], dtype=float),
+        stress=stress,
+    )
+    atoms.info.update({"oxide": oxide, "polymorph": polymorph, "xc": "PBE"})
+    return atoms
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input", required=True, type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--overwrite", action="store_true", help="Replace existing split databases"
+    )
+    args = parser.parse_args()
+
+    records = load_records(args.input)
+    if not records:
+        raise ValueError(f"No PBE EOS structures found in {args.input}")
+    args.output.mkdir(parents=True, exist_ok=True)
+    paths = {split: args.output / f"{split}.db" for split in ("train", "val", "test")}
+    existing = [str(path) for path in paths.values() if path.exists()]
+    if existing and not args.overwrite:
+        raise FileExistsError("Output exists; pass --overwrite to replace: " + ", ".join(existing))
+    for path in paths.values():
+        if path.exists():
+            path.unlink()
+
+    split_for = split_groups(records, args.seed)
+    counts = Counter()
+    databases = {split: connect(path) for split, path in paths.items()}
+    for oxide, polymorph, calculation in records:
+        split = split_for[(oxide, polymorph)]
+        databases[split].write(
+            to_atoms(oxide, polymorph, calculation),
+            oxide=oxide,
+            polymorph=polymorph,
+            xc="PBE",
+        )
+        counts[split] += 1
+
+    manifest = {
+        "source": str(args.input.resolve()),
+        "seed": args.seed,
+        "split_strategy": "oxide-polymorph grouped 80/10/10",
+        "stress_source_unit": "GPa",
+        "stress_output_unit": "eV/Angstrom^3",
+        "counts": dict(counts),
+        "groups": {f"{oxide}/{polymorph}": split for (oxide, polymorph), split in sorted(split_for.items())},
+    }
+    with (args.output / "manifest.json").open("w") as handle:
+        json.dump(manifest, handle, indent=2, sort_keys=True)
+    print(json.dumps(manifest["counts"], sort_keys=True))
+    print(f"saved dataset: {args.output.resolve()}")
+
+
+if __name__ == "__main__":
+    main()

relax.py

@@ -0,0 +1,62 @@
+"""Relax a periodic structure with an eSEN ASE calculator."""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+os.environ.setdefault(
+    "ONESCIENCE_ESEN_JD_PATH",
+    os.path.join(os.path.dirname(__file__), "weight", "Jd.pt"),
+)
+
+from ase.build import bulk
+from ase.filters import FrechetCellFilter
+from ase.io import read, write
+from ase.optimize import BFGS
+
+from onescience.utils.esen import eSENCalculator
+
+
+def default_checkpoint() -> str:
+    return os.path.join(os.path.dirname(__file__), "weight", "esen_30m_mptrj.pt")
+
+
+def load_structure(path: str | None):
+    if path:
+        return read(path)
+    return bulk("Si", "diamond", a=5.50)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", default=default_checkpoint())
+    parser.add_argument("--input", help="CIF, POSCAR, XYZ, or another ASE-readable structure")
+    parser.add_argument("--fmax", type=float, default=0.05)
+    parser.add_argument("--steps", type=int, default=100)
+    parser.add_argument("--output", default="relaxed.cif")
+    parser.add_argument(
+        "--fixed-cell",
+        action="store_true",
+        help="relax atomic positions only; by default the cell is relaxed too",
+    )
+    parser.add_argument("--device", default="cuda")
+    args = parser.parse_args()
+
+    atoms = load_structure(args.input)
+    if not args.fixed_cell and not atoms.pbc.all():
+        parser.error("cell relaxation requires periodic boundaries; use --fixed-cell")
+    atoms.calc = eSENCalculator.from_checkpoint(args.checkpoint, device=args.device)
+    target = atoms if args.fixed_cell else FrechetCellFilter(atoms)
+    optimizer = BFGS(target, logfile="relax.log", trajectory="relax.traj")
+    optimizer.run(fmax=args.fmax, steps=args.steps)
+    write(args.output, atoms)
+    print("formula:", atoms.get_chemical_formula())
+    print("atoms:", len(atoms))
+    print("steps:", optimizer.nsteps)
+    print("energy (eV):", atoms.get_potential_energy())
+    print("maximum force (eV/Angstrom):", max((atoms.get_forces() ** 2).sum(1) ** 0.5))
+
+
+if __name__ == "__main__":
+    main()

single_point.py

@@ -0,0 +1,47 @@
+"""Run one eSEN single-point calculation through ASE."""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+os.environ.setdefault(
+    "ONESCIENCE_ESEN_JD_PATH",
+    os.path.join(os.path.dirname(__file__), "weight", "Jd.pt"),
+)
+
+from ase.build import bulk
+from ase.io import read
+
+from onescience.utils.esen import eSENCalculator
+
+
+def default_checkpoint() -> str:
+    return os.path.join(os.path.dirname(__file__), "weight", "esen_30m_mptrj.pt")
+
+
+def load_structure(path: str | None):
+    if path:
+        return read(path)
+    return bulk("Si", "diamond", a=5.43)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", default=default_checkpoint())
+    parser.add_argument("--input", help="CIF, POSCAR, XYZ, or another ASE-readable structure")
+    parser.add_argument("--device", default="cuda")
+    args = parser.parse_args()
+
+    atoms = load_structure(args.input)
+    atoms.calc = eSENCalculator.from_checkpoint(args.checkpoint, device=args.device)
+
+    print("formula:", atoms.get_chemical_formula())
+    print("atoms:", len(atoms))
+    print("energy (eV):", atoms.get_potential_energy())
+    print("forces (eV/Angstrom):\n", atoms.get_forces())
+    print("stress (eV/Angstrom^3):", atoms.get_stress())
+
+
+if __name__ == "__main__":
+    main()

weight/Jd.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4059c45be246dcb6c49c545670b65c56550eb0c2e7a9c92b4b50a92d370dbe2
+size 21697