【MIIT program】SevenNet模型复现

interatomic_potentials/configs/sevennet/README.md

@@ -0,0 +1,103 @@
+# SevenNet - Interatomic Potential Model
+
+## 任务简介
+
+SevenNet 是一个基于图神经网络的机器学习原子间势函数模型，用于预测分子和材料的能量与力。
+
+## 模型简介
+
+SevenNet 模型结构包括：
+- **原子嵌入层**: 将原子序数映射为特征向量
+- **径向基函数 (RBF)**: 编码原子间距离信息
+- **消息传递块**: 多层图神经网络进行信息聚合
+- **能量预测头**: 预测原子能量并求和得到总能量
+
+## 环境依赖
+
+```bash
+# 安装 PaddlePaddle
+pip install paddlepaddle==3.3.1
+
+# 安装依赖
+pip install ase numpy tqdm pyyaml
+```
+
+## 数据准备
+
+数据集格式：extxyz 格式，包含原子坐标、能量和力标签。
+
+示例数据位于：`tests/data/systems/hfo2.extxyz`
+
+## 训练命令
+
+```bash
+# 使用默认配置训练
+cd interatomic_potentials
+python train.py --config configs/sevennet/sevennet_hfo2.yaml
+
+# 指定输出目录
+python train.py --config configs/sevennet/sevennet_hfo2.yaml --output-dir ./output/sevennet_hfo2
+```
+
+## 推理命令
+
+```bash
+cd interatomic_potentials
+python infer.py --model ./output/sevennet_hfo2/model_final.pdparams --structure tests/data/systems/hfo2.extxyz
+```
+
+## 配置说明
+
+### 模型配置 (`model`)
+- `type`: 模型类型，固定为 "sevennet"
+- `num_species`: 原子种类数，默认 100
+- `hidden_dim`: 隐藏层维度，默认 128
+- `num_message_layers`: 消息传递层数，默认 4
+- `num_rbf`: 径向基函数数量，默认 32
+- `cutoff`: 截断半径（Å），默认 5.0
+
+### 数据集配置 (`dataset`)
+- `type`: 数据集类型，固定为 "extxyz"
+- `path`: 数据集路径
+- `cutoff`: 截断半径
+- `valid_ratio`: 验证集比例，默认 0.1
+
+### 训练配置 (`trainer`)
+- `epoch`: 训练轮数
+- `per_epoch`: 每多少轮保存一次
+- `seed`: 随机种子
+- `device`: 设备，"auto"、"gpu" 或 "cpu"
+
+## 参考结果
+
+使用 `sevennet_hfo2.yaml` 配置训练 HfO2 数据集：
+- 训练轮数：2
+- 初始训练损失：~888
+- 最终验证损失：~794
+
+## 目录结构
+
+```
+interatomic_potentials/
+├── train.py              # 训练入口脚本
+├── infer.py              # 推理入口脚本
+└── configs/
+    └── sevennet/
+        ├── sevennet_hfo2.yaml        # 默认配置
+        ├── sevennet_hfo2_large.yaml  # 大型配置
+        └── README.md                 # 本文件
+```
+
+## 代码结构
+
+```
+ppmat/
+└── models/
+    └── sevennet/
+        ├── __init__.py
+        └── sevennet_model.py  # SevenNet 模型实现
+```
+
+## 参考
+
+本模型参考了 SevenNet 系列原子间势函数的设计理念。

interatomic_potentials/configs/sevennet/sevennet_hfo2.yaml

@@ -0,0 +1,40 @@
+model:
+  type: "sevennet"
+  num_species: 100
+  hidden_dim: 32
+  num_message_layers: 3
+  num_rbf: 32
+  cutoff: 4.0
+
+dataset:
+  type: "extxyz"
+  path: "/home/user/cv/zyy/PaddleSeveNnet/tests/data/systems/hfo2.extxyz"
+  cutoff: 4.0
+  valid_ratio: 0.1
+
+dataloader:
+  batch_size: 2
+  shuffle: true
+  num_workers: 0
+
+optimizer:
+  type: "adam"
+  lr: 0.001
+
+lr_scheduler:
+  type: "constant"
+
+loss:
+  type: "mse"
+  force_loss_weight: 0.1
+
+trainer:
+  epoch: 2
+  per_epoch: 1
+  seed: 42
+  device: "auto"
+
+log:
+  save_dir: "./output/sevennet_hfo2"
+  log_interval: 1
+  save_interval: 1

interatomic_potentials/configs/sevennet/sevennet_hfo2_large.yaml

@@ -0,0 +1,40 @@
+model:
+  type: "sevennet"
+  num_species: 100
+  hidden_dim: 128
+  num_message_layers: 5
+  num_rbf: 64
+  cutoff: 5.0
+
+dataset:
+  type: "extxyz"
+  path: "/home/user/cv/zyy/PaddleSeveNnet/tests/data/systems/hfo2.extxyz"
+  cutoff: 5.0
+  valid_ratio: 0.1
+
+dataloader:
+  batch_size: 4
+  shuffle: true
+  num_workers: 0
+
+optimizer:
+  type: "adam"
+  lr: 0.001
+
+lr_scheduler:
+  type: "constant"
+
+loss:
+  type: "mse"
+  force_loss_weight: 0.1
+
+trainer:
+  epoch: 10
+  per_epoch: 2
+  seed: 42
+  device: "auto"
+
+log:
+  save_dir: "./output/sevennet_hfo2_large"
+  log_interval: 1
+  save_interval: 2

interatomic_potentials/infer.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+PaddleMaterials - Interatomic Potentials Inference Script
+推理入口：python infer.py --model model.pdparams --structure structure.extxyz
+"""
+
+import argparse
+import os
+import sys
+import numpy as np
+import paddle
+from ase.io import read
+from ase.neighborlist import neighbor_list
+
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, PROJECT_ROOT)
+
+from ppmat.models.sevennet import PurePaddleSevenNet
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Infer interatomic potential model")
+    parser.add_argument("--model", type=str, required=True, help="Path to model checkpoint")
+    parser.add_argument("--structure", type=str, required=True, help="Path to structure file")
+    parser.add_argument("--cutoff", type=float, default=5.0, help="Cutoff radius")
+    return parser.parse_args()
+
+
+def atoms_to_graph(atoms, cutoff):
+    positions = atoms.get_positions().astype("float32")
+    numbers = atoms.get_atomic_numbers().astype("int64")
+    i, j = neighbor_list("ij", atoms, cutoff)
+    edge_index = np.vstack([i, j]).astype("int64")
+
+    return {
+        "z": paddle.to_tensor(numbers, dtype="int64"),
+        "pos": paddle.to_tensor(positions, dtype="float32"),
+        "edge_index": paddle.to_tensor(edge_index, dtype="int64"),
+    }
+
+
+def predict(model, atoms, cutoff, energy_mean, energy_std):
+    graph = atoms_to_graph(atoms, cutoff)
+    num_atoms = len(atoms)
+
+    with paddle.no_grad():
+        out = model(graph)
+        raw_total = float(out["total_energy"].numpy())
+        pred_total = raw_total * energy_std + energy_mean
+        pred_per_atom = pred_total / num_atoms
+
+    return pred_total, pred_per_atom
+
+
+def main():
+    args = parse_args()
+
+    # 加载模型
+    print(f"Loading model: {args.model}")
+    checkpoint = paddle.load(args.model)
+    state_dict = checkpoint["model_state_dict"]
+    energy_mean = float(checkpoint.get("energy_mean", 0.0))
+    energy_std = float(checkpoint.get("energy_std", 1.0))
+
+    # 获取配置
+    config = checkpoint.get("config", {})
+    model_cfg = config.get("model", {})
+
+    # 创建模型
+    model = PurePaddleSevenNet(
+        num_species=model_cfg.get("num_species", 100),
+        hidden_dim=model_cfg.get("hidden_dim", 128),
+        num_message_layers=model_cfg.get("num_message_layers", 4),
+        num_rbf=model_cfg.get("num_rbf", 32),
+        cutoff=model_cfg.get("cutoff", args.cutoff),
+    )
+    model.set_state_dict(state_dict)
+    model.eval()
+
+    print(f"Model loaded successfully")
+    print(f"Energy mean: {energy_mean:.6f}, std: {energy_std:.6f}")
+
+    # 读取结构
+    print(f"\nReading structure: {args.structure}")
+    atoms_list = read(args.structure, index=":")
+    if not isinstance(atoms_list, list):
+        atoms_list = [atoms_list]
+    print(f"Loaded {len(atoms_list)} structures")
+
+    # 预测
+    print("\n=== Prediction Results ===")
+    for idx, atoms in enumerate(atoms_list):
+        pred_total, pred_per_atom = predict(model, atoms, args.cutoff, energy_mean, energy_std)
+
+        ref_energy = None
+        for key in ["y_energy", "energy", "free_energy", "total_energy"]:
+            if key in atoms.info:
+                try:
+                    ref_energy = float(atoms.info[key])
+                    break
+                except Exception:
+                    pass
+
+        print(f"\nStructure {idx + 1}:")
+        print(f"  Atoms: {len(atoms)}")
+        print(f"  Predicted total energy: {pred_total:.6f} eV")
+        print(f"  Predicted per-atom energy: {pred_per_atom:.6f} eV/atom")
+
+        if ref_energy is not None:
+            print(f"  Reference energy: {ref_energy:.6f} eV")
+            print(f"  Error: {abs(pred_total - ref_energy):.6f} eV")
+
+    print("\nInference completed!")
+
+
+if __name__ == "__main__":
+    main()