File size: 1,591 Bytes
06c3567 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | from ase.build import bulk, molecule, fcc100, add_adsorbate
# 自动定位 UMA 旋转基文件 Jd.pt
import os
_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
if os.path.isfile(_JD_PATH):
os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
from onescience.utils.uma.units.mlip_unit import load_predict_unit
from onescience.datapipes.materials.custom_stack.core.atomic_data import AtomicData, atomicdata_list_to_batch
# 1. 创建异构结构
h2o = molecule("H2O")
h2o.info.update({"charge": 0, "spin": 1})
pt = bulk("Pt")
slab = fcc100("Cu", (3, 3, 3), vacuum=8, periodic=True)
adsorbate = molecule("CO")
add_adsorbate(slab, adsorbate, 2.0, "bridge")
# 2. 结构转为 AtomicData，并指定不同 task_name
atomic_data_list = [
AtomicData.from_ase(
h2o, task_name="omol", r_data_keys=["spin", "charge"], molecule_cell_size=12
),
AtomicData.from_ase(pt, task_name="omat"),
AtomicData.from_ase(slab, task_name="oc20"),
]
# 3. 合成 batch
batch = atomicdata_list_to_batch(atomic_data_list)
# 4. 加载 UMA 模型
predictor = load_predict_unit(
"../weight/uma-s-1p1_converted.pt", device="cuda"#替换为你的检查点路径
)
# 5. 执行联合推理
preds = predictor.predict(batch)
# 6. 输出每个结构结果
for i in range(len(preds["energy"])):
energy = preds["energy"][i].item()
forces = preds["forces"][batch.batch == i].cpu().numpy()
print(f"\n[Structure {i}]")
print("Predicted energy:", energy)
print("Predicted forces:\n", forces)
|