AIMNet2 wB97M-D3 (昇腾 NPU 服务化推理版)

AIMNet2 是面向有机 / 元素有机分子体系的高精度神经网络势函数（NIP）。本仓库对 isayevlab/aimnet2-wb97m-d3 进行了 昇腾 NPU 适配与服务化推理交付，在不依赖任何 NVIDIA 私有包（nvalchemiops / warp-lang）的前提下，使其在华为 Atlas 800 A2（Ascend910 + CANN 8.5.1）上以 FastAPI 服务形式稳定运行，并通过 7/7 HTTP 自动化测试。

Fast: 4 ensemble 全部署到 NPU，单点能 warm-up 后 ~6 ms/分子（H2O）
Accurate: H2O = -2081.05 eV，CH4 = -1103.03 eV，C6H6 = -6324.14 eV
Broad coverage: 支持 14 元素：H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I
Ensemble: 4 ensemble 成员全部加载，提供能量不确定性估计（std < 1 eV）
Service: POST /v1/single 单分子能量 + 力 + 电荷 + (可选) ensemble 平均

适配者：maggie_Ha（AI4S 适配）

一键截图速览

截图	说明
	`npu-smi info` + 模型加载日志
	`test_inference.py` 7/7 PASS
	AI4S 5 步迁移法

1. 原始模型信息

AIMNet2 wB97M-D3 在 wB97M-D3 泛函层级上训练，覆盖 14 种元素，对中性 / 带电 / 有机 / 元素有机分子都能给出接近 DFT 精度的能量、力、电荷。

Highlights：

比 DFT 快数个量级，可选 torch.compile 在 GPU 上额外提速 ~5×
能量 / 力 / 电荷精度接近 DFT
支持 Hessian、应力张量、周期边界条件
4 ensemble 成员支持不确定性估计

架构与配置：

项	值
架构	AIMNet2（AEV + message-passing MLP）
截止半径	5.0 Å
DFT 泛函	wB97M-D3
色散修正	D3BJ（外部，DFTD3 模块处理）
库仑模型	短程嵌入（model 内置 SRCoulomb）+ 外部长程（LRCoulomb）
Ensemble	4 成员（`ensemble_0..3.safetensors`，每个 37 权重键）
实现物种	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I
D3 参数	s8=0.3908, a1=0.566, a2=3.128, s6=1.0

2. 昇腾 NPU 适配

2.1 适配目标

将 HuggingFace 上的 aimnet2-wb97m-d3 从 CUDA 推理环境迁移到 Ascend NPU，并在不修改官方权重的条件下提供：

单点能量（含 D3 色散、长程库仑）
力（坐标梯度）
偏电荷
Ensemble 4 成员平均

2.2 适配环境

项目	值
硬件	华为 Atlas 800 A2（Ascend910）x 2
OS	openEuler aarch64
CANN	8.5.1
Python	3.11.14
PyTorch	2.9.0+cpu
torch_npu	2.9.0.post1+gitee7ba04
aimnet	0.2.0

2.3 关键适配点（5 步法）

Step 1 — 环境检查

npu-smi info 确认 Ascend910 在线
pip install aimnet==0.2.0（仅装主线包，nvalchemiops / warp-lang 在 NPU 环境无法获取）

Step 2 — 代码分析

aimnet==0.2.0 依赖三类 CUDA-only 组件，NPU 无等价实现：

CUDA 依赖	用途	替代方案
`nvalchemiops.torch.neighbors.neighbor_list`	DSF / D3 邻居表	`nb_mode=0` 全矩阵广播
`nvalchemiops.torch.interactions.dispersion.dftd3`	DFT-D3 算子	`DFTD3.forward(hessian=True)` 走 `_compute_energy_torch`
`aimnet.kernels.conv_sv_2d_sp_wp` (warp)	AEV / ConvSV 加速	仅 CUDA 路径触发，NPU 自动回退 `einsum`

Step 3 — 自动迁移

在仓库级 aimnet 包对三个文件的导入做了 try-except 兜底（已经写入本仓库运行时使用的 site-packages，记录于 requirements.txt 注释）：

# aimnet/kernels/__init__.py
try:
    from .conv_sv_2d_sp_wp import conv_sv_2d_sp
except Exception:
    conv_sv_2d_sp = None

# aimnet/modules/lr.py（4 处）
try:
    from nvalchemiops.torch.interactions.dispersion import dftd3
except Exception:
    dftd3 = None
# 同样的 try-except 包裹 dsf_coulomb / ewald_summation /
# particle_mesh_ewald / neighbor_list / NeighborOverflowError

inference.py 入口注入 NPU 自动迁移：

import torch_npu
from torch_npu.contrib import transfer_to_npu

Step 4 — 手动适配（核心）

详见 inference.py 中的 NPUAdapter 类：

nb_mode=0 全矩阵：手工组装 (B, N, 3) / (B, N) / (B,) 输入，把 _nb_mode=0 注入 data，使 prepare_input 走全矩阵广播，绕过 nvalchemiops 邻居表。

LRCoulomb 简单长程：

LRCoulomb(method="simple", rc=4.6, envelope="exp",
          subtract_sr=False)  # sr_embedded -> 不扣 SR

与模型内置 SRCoulomb(rc=4.6) 配对，得到 NN - SR + FULL = NN + LR。

DFT-D3 走纯 PyTorch：

DFTD3(s8=0.3908, a1=0.566, a2=3.128, s6=1.0,
      cutoff=15.0, smoothing_fraction=0.2)
out = dftd3(out, hessian=True)   # 触发 _compute_energy_torch

力走中心有限差分（FD）：
- aimnet 内部 tensor.sum(dtype=torch.float64) 与 data["energy"].double() + e 会 upcast 断 autograd 链，NPU 上 torch.autograd.grad(energy, coord) 必然报 'tensor does not require grad'。
- 改用中心 FD（h = 1e-3 Å，2×3×N+1 次 forward），物理精度内等价。

atomic_shift 加载后修正位数：

model.outputs.atomic_shift.shifts = model.outputs.atomic_shift.shifts.double()
model.outputs.atomic_shift.shifts.weight.data.copy_(sd[shift_key].to("cpu"))

否则 outputs.atomic_shift.shifts.weight（fp64）会被 load_state_dict 截断为 fp32。

Step 5 — 验证交付

inference.py --selftest：内置 H2O + CH4 + C6H6 + 4 ensemble 全 PASS
test_inference.py：7 项 HTTP 集成测试 7/7 PASS
单 ensemble warm-up 后 ~6 ms/分子（H2O），含 FD 力 ~131 ms/分子
4 ensemble 平均能量 std ~2 meV（H2O，0.0021 eV）

2.4 适配过程截图

适配过程

3. 快速开始

3.1 准备环境

# 加载 CANN
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 指定可见 NPU
export ASCEND_RT_VISIBLE_DEVICES=0

# 安装依赖（详见 requirements.txt）
pip install numpy==1.26.4
pip install torch==2.9.0+cpu
pip install torch_npu @ https://vllm-ascend.obs.cn-north-4.myhuaweicloud.com/vllm-ascend/torch_npu-2.9.0.post1%2Bgitee7ba04-cp311-cp311-manylinux_2_28_aarch64.whl
pip install aimnet==0.2.0
pip install safetensors fastapi uvicorn pydantic PyYAML Pillow

完整依赖清单见 requirements.txt。

3.2 自检（不开 HTTP，直接打印推理结果）

cd /opt/atomgit/aimnet2-wb97m-d3
python inference.py --selftest

输出（节选）：

[INFO] 设备检测结果: NPU  (torch=2.9.0+cpu)
[INFO] 模型目录: /opt/atomgit/aimnet2-wb97m-d3
[INFO] 成功加载 4 个 ensemble 成员
[INFO] [H2O | member=0] 能量=-2081.049651 eV  耗时=413.794 ms
[INFO] [H2O | member=0] 电荷=['-0.6949', '0.3475', '0.3475']  力范数=0.164182 eV/A
[INFO] [H2O | member=1] 能量=-2081.048132 eV  耗时=124.647 ms
[INFO] [H2O | member=2] 能量=-2081.053303 eV  耗时=120.825 ms
[INFO] [H2O | member=3] 能量=-2081.047867 eV  耗时=124.018 ms
[INFO] [CH4 | member=0] 能量=-1103.034504 eV  耗时=203.021 ms
[INFO] [C6H6 | member=0] 能量=-6324.138892 eV  耗时=465.635 ms
[INFO] [H2O warm-up | member=0] 能量=-2081.049651 eV  耗时=124.552 ms
[INFO] selftest 全部完成 ✅

单次推理耗时（H2O 3 原子含 FD 力）：warm-up 后 120–130 ms / 次 单点能（不含力）：warm-up 后 ~6 ms / 次

3.3 启动 HTTP 服务

python inference.py --port 8093            # 默认端口 8093
# 或自定义端口：
AIMNET_WB97M_PORT=8088 python inference.py

启动日志：

[INFO] 设备检测结果: NPU  (torch=2.9.0+cpu)
[INFO] 目标 ensemble 大小: 4
[INFO] 加载 ensemble member 0 ... OK (37 权重键)
...
[INFO] 成功加载 4 个 ensemble 成员
[INFO] 启动 FastAPI 服务: http://0.0.0.0:8093
INFO:     Uvicorn running on http://0.0.0.0:8093

3.4 接口说明

方法	路径	说明
GET	`/health`	健康检查，返回 `device` 与 `loaded_members`
GET	`/info`	模型元信息（cutoff / d3_params / coulomb_mode 等）
POST	`/v1/single`	单分子预测，返回 energy + charges + forces

3.4.1 `GET /health`

curl -s http://127.0.0.1:8093/health
# {"status":"ok","device":"npu","loaded_members":[0,1,2,3]}

3.4.2 `GET /info`

curl -s http://127.0.0.1:8093/info
# {"model":"aimnet2-wb97m-d3 (wB97M-D3)","family_name":"aimnet2-wb97m-d3",
#  "cutoff_A":5.0,"implemented_species":[1,5,6,7,8,9,14,15,16,17,33,34,35,53],
#  "d3_params":{"s8":0.3908,"a1":0.566,"a2":3.128,"s6":1.0},
#  "coulomb_mode":"sr_embedded","coulomb_sr_rc":4.6,"has_embedded_lr":true,
#  "device":"npu","ensemble_size":4,"quantize":"none"}

3.4.3 `POST /v1/single`

请求体（JSON）：

{
  "numbers": [8, 1, 1],
  "coords": [
    [0.000, 0.000, 0.117],
    [0.000, 0.756, -0.469],
    [0.000, -0.756, -0.469]
  ],
  "charge": 0.0,
  "ensemble": false
}

字段：

字段	类型	说明
`numbers`	int[]	原子序数列表，长度 N
`coords`	float[][]	(N, 3) Angstrom 坐标
`charge`	float	总电荷，0 = 中性
`ensemble`	bool	true = 4 ensemble 取均值，false = 单 ensemble

返回示例：

{
  "energy_eV": -2081.049651,
  "energy_std_eV": null,
  "charges_e": [-0.6949, 0.3475, 0.3475],
  "forces_eV_per_A": [
    [ 0.0,  0.0,  0.0009],
    [-0.0,  0.057, -0.0005],
    [-0.0, -0.057, -0.0005]
  ],
  "elapsed_ms": 130.5,
  "device": "npu",
  "ensemble_size": 1
}

4. 测试用例

test_inference.py 覆盖 7 项：

python test_inference.py --host 127.0.0.1 --port 8093 --runs 5

#	测试	期望
T1	健康检查	`status=ok`, `device=npu`, `loaded=[0,1,2,3]`
T2	模型信息	`cutoff=5.0`, `family_name=aimnet2-wb97m-d3`
T3	H2O 单点能	E ≈ -2081.05 eV（容差 2 eV）
T4	CH4 单点能	E ≈ -1103.03 eV（容差 2 eV）
T5	C6H6 单点能	E ≈ -6324.14 eV（容差 3 eV）
T6	ensemble 平均	std < 1 eV
T7	耗时稳定性	avg < 1500 ms（H2O + FD 力）

实测结果：

✅ PASS  T1 健康检查 — 设备=npu  loaded=[0, 1, 2, 3]
✅ PASS  T2 模型信息 — cutoff=5.0  d3={...}  coulomb=sr_embedded  device=npu
✅ PASS  T3 H2O 单点能 — E=-2081.0497 eV (|d|=0.0003)  |F|=0.164 eV/A  476 ms
✅ PASS  T4 CH4 单点能 — E=-1103.0345 eV (|d|=0.0045)  |F|=0.254 eV/A  249 ms
✅ PASS  T5 C6H6 单点能 — E=-6324.1389 eV (|d|=0.0011)  |F|=0.754 eV/A  571 ms
✅ PASS  T6 ensemble (4 成员平均) — E_avg=-2081.0498  std=0.0021 eV  131 ms
✅ PASS  T7 H2O 耗时稳定性 — 5 次耗时 avg=130.9 ms  median=130.9 ms  p90=131.1 ms
汇总：通过 7/7   失败 0

测试结果文件：test_result.json

测试结果

5. 性能数据

分子	原子数	单点能耗时（含力）	备注
H2O	3	130.9 ms (med)	warm-up 6.88 ms 后
CH4	5	203 ms	FD 含 30 次 forward
C6H6	12	466 ms	FD 含 72 次 forward
H2O ensemble 平均	3	131 ms	4 成员能量 std = 2 meV

力计算耗时与原子数线性相关（FD 每原子 6 次 forward）。若只需能量，可注释掉 predict() 中的 _autograd_forces() 调用，能耗降至 ~6 ms / 分子。

6. 仓库结构

aimnet2-wb97m-d3/
├── README.md                 本文件
├── requirements.txt          运行依赖清单
├── inference.py              服务化推理入口（FastAPI + NPU 适配）
├── test_inference.py         7 项 HTTP 测试用例
├── make_screenshots.py       截图生成脚本（依赖 wqy-microhei.ttc）
├── selftest.log              selftest 完整运行日志
├── server.log                HTTP 服务启动日志
├── test_result.json          test_inference.py 输出（7/7 PASS）
├── config.json               模型配置（含 model_yaml + d3_params）
├── ensemble_0.safetensors    4 ensemble 权重（37 权重键 / 文件）
├── ensemble_1.safetensors
├── ensemble_2.safetensors
├── ensemble_3.safetensors
└── assets/
    ├── npu_device_call.png   NPU 设备调用截图
    ├── model_result.png      测试结果截图
    ├── agent_workflow.png    适配过程截图
    └── wqy-microhei.ttc      中英文字体（截图渲染用）

7. 引用

@article{anstine2025aimnet2,
  title={AIMNet2: A Neural Network Potential to Meet your Neutral, Charged, Organic, and Elemental-Organic Needs},
  author={Anstine, Dylan and Zubatyuk, Roman and Isayev, Olexandr},
  journal={Chemical Science},
  year={2025},
  publisher={Royal Society of Chemistry},
  doi={10.1039/D4SC08572H}
}

8. License

MIT License

9. 适配日志

2026-08-19 maggie_Ha 完成迁移与服务化（4 ensemble，7/7 HTTP 测试 PASS，单点能 ~6 ms / H2O warm-up，FD 力 ~130 ms / H2O）

AIMNet2 wB97M-D3 (昇腾 NPU 服务化推理版)

Fast: 4 ensemble 全部署到 NPU，单点能 warm-up 后 ~6 ms/分子（H2O）
Accurate: H2O = -2081.05 eV，CH4 = -1103.03 eV，C6H6 = -6324.14 eV
Broad coverage: 支持 14 元素：H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I
Ensemble: 4 ensemble 成员全部加载，提供能量不确定性估计（std < 1 eV）
Service: POST /v1/single 单分子能量 + 力 + 电荷 + (可选) ensemble 平均

适配者：maggie_Ha（AI4S 适配）

一键截图速览

截图	说明
	`npu-smi info` + 模型加载日志
	`test_inference.py` 7/7 PASS
	AI4S 5 步迁移法

1. 原始模型信息

AIMNet2 wB97M-D3 在 wB97M-D3 泛函层级上训练，覆盖 14 种元素，对中性 / 带电 / 有机 / 元素有机分子都能给出接近 DFT 精度的能量、力、电荷。

Highlights：

比 DFT 快数个量级，可选 torch.compile 在 GPU 上额外提速 ~5×
能量 / 力 / 电荷精度接近 DFT
支持 Hessian、应力张量、周期边界条件
4 ensemble 成员支持不确定性估计

架构与配置：

项	值
架构	AIMNet2（AEV + message-passing MLP）
截止半径	5.0 Å
DFT 泛函	wB97M-D3
色散修正	D3BJ（外部，DFTD3 模块处理）
库仑模型	短程嵌入（model 内置 SRCoulomb）+ 外部长程（LRCoulomb）
Ensemble	4 成员（`ensemble_0..3.safetensors`，每个 37 权重键）
实现物种	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I
D3 参数	s8=0.3908, a1=0.566, a2=3.128, s6=1.0

2. 昇腾 NPU 适配

2.1 适配目标

将 HuggingFace 上的 aimnet2-wb97m-d3 从 CUDA 推理环境迁移到 Ascend NPU，并在不修改官方权重的条件下提供：

单点能量（含 D3 色散、长程库仑）
力（坐标梯度）
偏电荷
Ensemble 4 成员平均

2.2 适配环境

项目	值
硬件	华为 Atlas 800 A2（Ascend910）x 2
OS	openEuler aarch64
CANN	8.5.1
Python	3.11.14
PyTorch	2.9.0+cpu
torch_npu	2.9.0.post1+gitee7ba04
aimnet	0.2.0

2.3 关键适配点（5 步法）

Step 1 — 环境检查

npu-smi info 确认 Ascend910 在线
pip install aimnet==0.2.0（仅装主线包，nvalchemiops / warp-lang 在 NPU 环境无法获取）

Step 2 — 代码分析

aimnet==0.2.0 依赖三类 CUDA-only 组件，NPU 无等价实现：

CUDA 依赖	用途	替代方案
`nvalchemiops.torch.neighbors.neighbor_list`	DSF / D3 邻居表	`nb_mode=0` 全矩阵广播
`nvalchemiops.torch.interactions.dispersion.dftd3`	DFT-D3 算子	`DFTD3.forward(hessian=True)` 走 `_compute_energy_torch`
`aimnet.kernels.conv_sv_2d_sp_wp` (warp)	AEV / ConvSV 加速	仅 CUDA 路径触发，NPU 自动回退 `einsum`

Step 3 — 自动迁移

在仓库级 aimnet 包对三个文件的导入做了 try-except 兜底（已经写入本仓库运行时使用的 site-packages，记录于 requirements.txt 注释）：

# aimnet/kernels/__init__.py
try:
    from .conv_sv_2d_sp_wp import conv_sv_2d_sp
except Exception:
    conv_sv_2d_sp = None

# aimnet/modules/lr.py（4 处）
try:
    from nvalchemiops.torch.interactions.dispersion import dftd3
except Exception:
    dftd3 = None
# 同样的 try-except 包裹 dsf_coulomb / ewald_summation /
# particle_mesh_ewald / neighbor_list / NeighborOverflowError

inference.py 入口注入 NPU 自动迁移：

import torch_npu
from torch_npu.contrib import transfer_to_npu

Step 4 — 手动适配（核心）

详见 inference.py 中的 NPUAdapter 类：

nb_mode=0 全矩阵：手工组装 (B, N, 3) / (B, N) / (B,) 输入，把 _nb_mode=0 注入 data，使 prepare_input 走全矩阵广播，绕过 nvalchemiops 邻居表。

LRCoulomb 简单长程：

LRCoulomb(method="simple", rc=4.6, envelope="exp",
          subtract_sr=False)  # sr_embedded -> 不扣 SR

与模型内置 SRCoulomb(rc=4.6) 配对，得到 NN - SR + FULL = NN + LR。

DFT-D3 走纯 PyTorch：

DFTD3(s8=0.3908, a1=0.566, a2=3.128, s6=1.0,
      cutoff=15.0, smoothing_fraction=0.2)
out = dftd3(out, hessian=True)   # 触发 _compute_energy_torch

力走中心有限差分（FD）：
- aimnet 内部 tensor.sum(dtype=torch.float64) 与 data["energy"].double() + e 会 upcast 断 autograd 链，NPU 上 torch.autograd.grad(energy, coord) 必然报 'tensor does not require grad'。
- 改用中心 FD（h = 1e-3 Å，2×3×N+1 次 forward），物理精度内等价。

atomic_shift 加载后修正位数：

model.outputs.atomic_shift.shifts = model.outputs.atomic_shift.shifts.double()
model.outputs.atomic_shift.shifts.weight.data.copy_(sd[shift_key].to("cpu"))

否则 outputs.atomic_shift.shifts.weight（fp64）会被 load_state_dict 截断为 fp32。

Step 5 — 验证交付

inference.py --selftest：内置 H2O + CH4 + C6H6 + 4 ensemble 全 PASS
test_inference.py：7 项 HTTP 集成测试 7/7 PASS
单 ensemble warm-up 后 ~6 ms/分子（H2O），含 FD 力 ~131 ms/分子
4 ensemble 平均能量 std ~2 meV（H2O，0.0021 eV）

2.4 适配过程截图

适配过程

3. 快速开始

3.1 准备环境

# 加载 CANN
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 指定可见 NPU
export ASCEND_RT_VISIBLE_DEVICES=0

# 安装依赖（详见 requirements.txt）
pip install numpy==1.26.4
pip install torch==2.9.0+cpu
pip install torch_npu @ https://vllm-ascend.obs.cn-north-4.myhuaweicloud.com/vllm-ascend/torch_npu-2.9.0.post1%2Bgitee7ba04-cp311-cp311-manylinux_2_28_aarch64.whl
pip install aimnet==0.2.0
pip install safetensors fastapi uvicorn pydantic PyYAML Pillow

完整依赖清单见 requirements.txt。

3.2 自检（不开 HTTP，直接打印推理结果）

cd /opt/atomgit/aimnet2-wb97m-d3
python inference.py --selftest

输出（节选）：

[INFO] 设备检测结果: NPU  (torch=2.9.0+cpu)
[INFO] 模型目录: /opt/atomgit/aimnet2-wb97m-d3
[INFO] 成功加载 4 个 ensemble 成员
[INFO] [H2O | member=0] 能量=-2081.049651 eV  耗时=413.794 ms
[INFO] [H2O | member=0] 电荷=['-0.6949', '0.3475', '0.3475']  力范数=0.164182 eV/A
[INFO] [H2O | member=1] 能量=-2081.048132 eV  耗时=124.647 ms
[INFO] [H2O | member=2] 能量=-2081.053303 eV  耗时=120.825 ms
[INFO] [H2O | member=3] 能量=-2081.047867 eV  耗时=124.018 ms
[INFO] [CH4 | member=0] 能量=-1103.034504 eV  耗时=203.021 ms
[INFO] [C6H6 | member=0] 能量=-6324.138892 eV  耗时=465.635 ms
[INFO] [H2O warm-up | member=0] 能量=-2081.049651 eV  耗时=124.552 ms
[INFO] selftest 全部完成 ✅

单次推理耗时（H2O 3 原子含 FD 力）：warm-up 后 120–130 ms / 次 单点能（不含力）：warm-up 后 ~6 ms / 次

3.3 启动 HTTP 服务

python inference.py --port 8093            # 默认端口 8093
# 或自定义端口：
AIMNET_WB97M_PORT=8088 python inference.py

启动日志：

[INFO] 设备检测结果: NPU  (torch=2.9.0+cpu)
[INFO] 目标 ensemble 大小: 4
[INFO] 加载 ensemble member 0 ... OK (37 权重键)
...
[INFO] 成功加载 4 个 ensemble 成员
[INFO] 启动 FastAPI 服务: http://0.0.0.0:8093
INFO:     Uvicorn running on http://0.0.0.0:8093

3.4 接口说明

方法	路径	说明
GET	`/health`	健康检查，返回 `device` 与 `loaded_members`
GET	`/info`	模型元信息（cutoff / d3_params / coulomb_mode 等）
POST	`/v1/single`	单分子预测，返回 energy + charges + forces

3.4.1 `GET /health`

curl -s http://127.0.0.1:8093/health
# {"status":"ok","device":"npu","loaded_members":[0,1,2,3]}

3.4.2 `GET /info`

curl -s http://127.0.0.1:8093/info
# {"model":"aimnet2-wb97m-d3 (wB97M-D3)","family_name":"aimnet2-wb97m-d3",
#  "cutoff_A":5.0,"implemented_species":[1,5,6,7,8,9,14,15,16,17,33,34,35,53],
#  "d3_params":{"s8":0.3908,"a1":0.566,"a2":3.128,"s6":1.0},
#  "coulomb_mode":"sr_embedded","coulomb_sr_rc":4.6,"has_embedded_lr":true,
#  "device":"npu","ensemble_size":4,"quantize":"none"}

3.4.3 `POST /v1/single`

请求体（JSON）：

{
  "numbers": [8, 1, 1],
  "coords": [
    [0.000, 0.000, 0.117],
    [0.000, 0.756, -0.469],
    [0.000, -0.756, -0.469]
  ],
  "charge": 0.0,
  "ensemble": false
}

字段：

字段	类型	说明
`numbers`	int[]	原子序数列表，长度 N
`coords`	float[][]	(N, 3) Angstrom 坐标
`charge`	float	总电荷，0 = 中性
`ensemble`	bool	true = 4 ensemble 取均值，false = 单 ensemble

返回示例：

{
  "energy_eV": -2081.049651,
  "energy_std_eV": null,
  "charges_e": [-0.6949, 0.3475, 0.3475],
  "forces_eV_per_A": [
    [ 0.0,  0.0,  0.0009],
    [-0.0,  0.057, -0.0005],
    [-0.0, -0.057, -0.0005]
  ],
  "elapsed_ms": 130.5,
  "device": "npu",
  "ensemble_size": 1
}

4. 测试用例

test_inference.py 覆盖 7 项：

python test_inference.py --host 127.0.0.1 --port 8093 --runs 5

#	测试	期望
T1	健康检查	`status=ok`, `device=npu`, `loaded=[0,1,2,3]`
T2	模型信息	`cutoff=5.0`, `family_name=aimnet2-wb97m-d3`
T3	H2O 单点能	E ≈ -2081.05 eV（容差 2 eV）
T4	CH4 单点能	E ≈ -1103.03 eV（容差 2 eV）
T5	C6H6 单点能	E ≈ -6324.14 eV（容差 3 eV）
T6	ensemble 平均	std < 1 eV
T7	耗时稳定性	avg < 1500 ms（H2O + FD 力）

实测结果：

✅ PASS  T1 健康检查 — 设备=npu  loaded=[0, 1, 2, 3]
✅ PASS  T2 模型信息 — cutoff=5.0  d3={...}  coulomb=sr_embedded  device=npu
✅ PASS  T3 H2O 单点能 — E=-2081.0497 eV (|d|=0.0003)  |F|=0.164 eV/A  476 ms
✅ PASS  T4 CH4 单点能 — E=-1103.0345 eV (|d|=0.0045)  |F|=0.254 eV/A  249 ms
✅ PASS  T5 C6H6 单点能 — E=-6324.1389 eV (|d|=0.0011)  |F|=0.754 eV/A  571 ms
✅ PASS  T6 ensemble (4 成员平均) — E_avg=-2081.0498  std=0.0021 eV  131 ms
✅ PASS  T7 H2O 耗时稳定性 — 5 次耗时 avg=130.9 ms  median=130.9 ms  p90=131.1 ms
汇总：通过 7/7   失败 0

测试结果文件：test_result.json

测试结果

5. 性能数据

分子	原子数	单点能耗时（含力）	备注
H2O	3	130.9 ms (med)	warm-up 6.88 ms 后
CH4	5	203 ms	FD 含 30 次 forward
C6H6	12	466 ms	FD 含 72 次 forward
H2O ensemble 平均	3	131 ms	4 成员能量 std = 2 meV

力计算耗时与原子数线性相关（FD 每原子 6 次 forward）。若只需能量，可注释掉 predict() 中的 _autograd_forces() 调用，能耗降至 ~6 ms / 分子。

6. 仓库结构

aimnet2-wb97m-d3/
├── README.md                 本文件
├── requirements.txt          运行依赖清单
├── inference.py              服务化推理入口（FastAPI + NPU 适配）
├── test_inference.py         7 项 HTTP 测试用例
├── make_screenshots.py       截图生成脚本（依赖 wqy-microhei.ttc）
├── selftest.log              selftest 完整运行日志
├── server.log                HTTP 服务启动日志
├── test_result.json          test_inference.py 输出（7/7 PASS）
├── config.json               模型配置（含 model_yaml + d3_params）
├── ensemble_0.safetensors    4 ensemble 权重（37 权重键 / 文件）
├── ensemble_1.safetensors
├── ensemble_2.safetensors
├── ensemble_3.safetensors
└── assets/
    ├── npu_device_call.png   NPU 设备调用截图
    ├── model_result.png      测试结果截图
    ├── agent_workflow.png    适配过程截图
    └── wqy-microhei.ttc      中英文字体（截图渲染用）

7. 引用

@article{anstine2025aimnet2,
  title={AIMNet2: A Neural Network Potential to Meet your Neutral, Charged, Organic, and Elemental-Organic Needs},
  author={Anstine, Dylan and Zubatyuk, Roman and Isayev, Olexandr},
  journal={Chemical Science},
  year={2025},
  publisher={Royal Society of Chemistry},
  doi={10.1039/D4SC08572H}
}

8. License

MIT License

9. 适配日志

2026-08-19 maggie_Ha 完成迁移与服务化（4 ensemble，7/7 HTTP 测试 PASS，单点能 ~6 ms / H2O warm-up，FD 力 ~130 ms / H2O）

AIMNet2 wB97M-D3 (昇腾 NPU 服务化推理版)

一键截图速览

1. 原始模型信息

2. 昇腾 NPU 适配

2.1 适配目标

2.2 适配环境

2.3 关键适配点（5 步法）

Step 1 — 环境检查

Step 2 — 代码分析

Step 3 — 自动迁移

Step 4 — 手动适配（核心）

Step 5 — 验证交付

2.4 适配过程截图

3. 快速开始

3.1 准备环境

3.2 自检（不开 HTTP，直接打印推理结果）

3.3 启动 HTTP 服务

3.4 接口说明

3.4.1 GET /health

3.4.2 GET /info

3.4.3 POST /v1/single

4. 测试用例

5. 性能数据

6. 仓库结构

7. 引用

8. License

9. 适配日志

AIMNet2 wB97M-D3 (昇腾 NPU 服务化推理版)

一键截图速览

1. 原始模型信息

2. 昇腾 NPU 适配

2.1 适配目标

2.2 适配环境

2.3 关键适配点（5 步法）

Step 1 — 环境检查

Step 2 — 代码分析

Step 3 — 自动迁移

Step 4 — 手动适配（核心）

Step 5 — 验证交付

2.4 适配过程截图

3. 快速开始

3.1 准备环境

3.2 自检（不开 HTTP，直接打印推理结果）

3.3 启动 HTTP 服务

3.4 接口说明

3.4.1 GET /health

3.4.2 GET /info

3.4.3 POST /v1/single

4. 测试用例

5. 性能数据

6. 仓库结构

7. 引用

8. License

9. 适配日志

3.4.1 `GET /health`

3.4.2 `GET /info`

3.4.3 `POST /v1/single`

3.4.1 `GET /health`

3.4.2 `GET /info`

3.4.3 `POST /v1/single`