Uh oh!
There was an error while loading. Please reload this page.
【MIIT program】Feature add SGEquiDiff - #300
Open
learncat163 wants to merge 15 commits into
Open
Conversation
learncat163
commented
Jun 24, 2026
Author
SGEquiDiff 迁移1. 概述将 SGEquiDiff（Symmetry-Guided Equivariant Diffusion）晶体生成模型从 PyTorch 框架迁移至 Paddle。 SGEquiDiff 是一种基于空间群对称性约束的等变扩散模型，用于生成符合晶体学对称性要求的晶体结构。该模型涉及复杂的对称群操作、非对称单元（ASU）坐标变换、Wyckoff 位置采样等计算。 项目的运行、训练、采样，在星河项目：https://aistudio.baidu.com/projectdetail/10304319 2. 权重信息2.1 原版权重原始的权重在：https://drive.google.com/drive/folders/1ONwO53i6oG1_yBqP0zPQ-IV_zLoIWyQR 如果不方便访问google drive，可以使用aistudio上的原版搬运镜像：https://aistudio.baidu.com/modelsdetail/47157/space 2.2 转换权重基于2种数据集 MP-20 数据集 和 MPTS-52 数据集 的权重，分别进行转换，地址：https://aistudio.baidu.com/modelsdetail/47158?modelId=47158 其中，paddle的http在线下载地址内置在 3. 精度对齐验证3.1 前向精度对齐前向精度对齐主要考虑4部分：
对比 PT/PD 4个子模块的前向传播输出，针对确定性输出（diffusion, space_group, lattice 非随机部分）计算最大绝对误差和平均绝对误差。 3.1.1 扩散模型（Diffusion）对齐
3.1.2 空间群模型（log_prob）
3.1.3 晶格模型（Lattice）
3.1.4 Wyckoff 模型Wyckoff 深度依赖随机数发生器RNG，导致难易对齐，除非对PT的原版代码直接进行hook修改，这有一定的破坏性，暂时无法实现。 前向精度对齐命令Python脚本,点击后展开脚本#!/usr/bin/env python3"""前向对齐验证脚本：对比 PT/PD 的各子模块前向传播输出。验证模块: 1. diffusion: predict_equivariant_vectors 2. space_group: log_prob 3. lattice: forward (lengths, angles, log_pfs, regularizer) 4. wyckoff: sample_and_log_prob用法: # PD 侧 python forward_align.py pd --output tmp/sgequidiff_diff_test/pd_forward.json # PT 侧 (在 sgequidiff 目录下运行) python forward_align.py pt --output tmp/sgequidiff_diff_test/pt_forward.json # 对比 python forward_align.py diff --pt tmp/sgequidiff_diff_test/pt_forward.json --pd tmp/sgequidiff_diff_test/pd_forward.json"""
import argparse
import json
import sys
from pathlib import Path
FIXED_SAMPLE = {
"space_group_number": 139,
"lattice_lengths": [5.141545295715332, 5.141545295715332, 9.469661712646484],
"lattice_angles": [90.0, 90.0, 90.0],
"element_indices": [64, 26, 13, 13],
"wyckoff_indices": [0, 5, 8, 9],
"frac_coords": [
[0.0000, 0.0000, 0.0000],
[0.2500, 0.2500, 0.2500],
[0.0000, 0.3354, 0.0000],
[0.2688, 0.5000, 0.0000],
],
}
TIMESTEPS = [100, 200, 500, 800]
def run_pd(args):
import numpy as np
import paddle
import yaml
import os
PROJECT_ROOT = Path(os.getcwd())
sys.path.insert(0, str(PROJECT_ROOT))
from ppmat.models.sgequidiff.diffusion_model import EquivariantDiffusionModelConfig
from ppmat.models.sgequidiff.crystal_sampler import CrystalSampler, CrystalSamplerConfig
from ppmat.models.sgequidiff.lattice_sampler import LatticeSamplerConfig
from ppmat.models.sgequidiff.wyckoff_transformer import WyckoffElementTransformerConfig
from ppmat.models.sgequidiff.non_equivariant_drift_modules import GNNConfig
from ppmat.models.sgequidiff.embedding_utils import set_global_embedding_tools
from ppmat.models.sgequidiff.data_utils import lattice_params_to_matrix_paddle
paddle.seed(42)
np.random.seed(42)
ckpt_dir = PROJECT_ROOT / "pretrained-pd/mp_20"# Load config from config.yaml
config_yaml_path = ckpt_dir / "config.yaml"
yaml.add_constructor(
'tag:yaml.org,2002:python/object/apply:pathlib.PosixPath',
lambda loader, node: str(loader.construct_sequence(node)[0]),
Loader=yaml.FullLoader,
)
with open(config_yaml_path, "r") as f:
raw_config = yaml.load(f, Loader=yaml.FullLoader)
# Initialize embeddings (required before model creation)
emb_config = raw_config.get("embeddings", {})
set_global_embedding_tools(
element_embedding_json_path=emb_config.get("element_embedding_json_path", "cgcnn_atom_init.json"),
space_group_embedding_json_path=emb_config.get("space_group_embedding_json_path", "init_tokens/space_group_features/space_group_embeddings_62dim.json"),
wyckoff_embedding_json_path=emb_config.get("wyckoff_embedding_json_path", "init_tokens/wyckoff_features/wyckoff_embeddings_231dim.json"),
chemistry_embedding_type=emb_config.get("chemistry_embedding_type", "identity"),
)
# Build model configs
model_config = raw_config.get("model", {})
frac_coord_config = model_config.get("frac_coord_config", {})
gnn_config_raw = frac_coord_config.get("gnn_config", {})
gnn_cfg = GNNConfig(
num_plane_wave_freqs=gnn_config_raw.get("num_plane_wave_freqs", 64),
num_cartesian_distance_gaussians=gnn_config_raw.get("num_cartesian_distance_gaussians", 64),
edge_hidden_dim=gnn_config_raw.get("edge_hidden_dim", 256),
atom_hidden_dim=gnn_config_raw.get("atom_hidden_dim", 256),
use_vpa=gnn_config_raw.get("use_vpa", True),
use_graph_norm=gnn_config_raw.get("use_graph_norm", True),
num_msg_pass_steps=gnn_config_raw.get("num_msg_pass_steps", 5),
cutoff=gnn_config_raw.get("cutoff", 7.0),
use_frac_coords_in_node_emb=gnn_config_raw.get("use_frac_coords_in_node_emb", False),
dataset_name=gnn_config_raw.get("dataset_name", "mp_20"),
)
diffusion_cfg = EquivariantDiffusionModelConfig(
model_type=frac_coord_config.get("model_type", "gnn"),
num_timesteps=1000,
noise_scheduler_num_monte_carlo_samples=frac_coord_config.get("noise_scheduler_num_monte_carlo_samples", 2500),
num_wn_lattice_translations=frac_coord_config.get("num_wn_lattice_translations", 3),
sigma_min=frac_coord_config.get("sigma_min", 0.002),
time_emb_dim=frac_coord_config.get("time_emb_dim", 128),
subsample_group_operations=frac_coord_config.get("subsample_group_operations", False),
gnn_config=gnn_cfg,
)
lattice_config_raw = model_config.get("lattice_config", {})
lattice_cfg = LatticeSamplerConfig(
input_dimension=lattice_config_raw.get("input_dimension", 128),
hidden_dimension=lattice_config_raw.get("hidden_dimension", 256),
num_hidden_layers=lattice_config_raw.get("n_emb_layers", 2),
min_lattice_length=lattice_config_raw.get("min_lattice_length", 2.0),
max_lattice_length=lattice_config_raw.get("max_lattice_length", 133.0),
min_lattice_angle=lattice_config_raw.get("min_lattice_angle", 60.0),
max_lattice_angle=lattice_config_raw.get("max_lattice_angle", 135.0),
lattice_param_dim=lattice_config_raw.get("lattice_param_dim", 32),
n_emb_layers=lattice_config_raw.get("n_emb_layers", 2),
lattice_length_bin_embedder_fourier_scale=lattice_config_raw.get("lattice_length_bin_embedder_fourier_scale", 2.0),
lattice_angle_bin_embedder_fourier_scale=lattice_config_raw.get("lattice_angle_bin_embedder_fourier_scale", 1.0),
lattice_length_embedder_fourier_scale=lattice_config_raw.get("lattice_length_embedder_fourier_scale", 5.0),
lattice_angle_embedder_fourier_scale=lattice_config_raw.get("lattice_angle_embedder_fourier_scale", 1.0),
)
we_config_raw = model_config.get("wyckoff_element_config", {})
we_cfg = WyckoffElementTransformerConfig(
hidden_dim=we_config_raw.get("hidden_dim", 256),
dataset_name=we_config_raw.get("dataset_name", "mp_20"),
num_heads=we_config_raw.get("num_heads", 2),
num_hidden_layers=we_config_raw.get("num_hidden_layers", 4),
dropout_rate=we_config_raw.get("dropout_rate", 0.1),
)
sampler_cfg = CrystalSamplerConfig(
diffusion_model_config=diffusion_cfg,
lattice_model_config=lattice_cfg,
transformer_config=we_cfg,
)
# Create CrystalSampler (contains all 4 sub-modules)
model = CrystalSampler(sampler_cfg)
model.eval()
# Load all 4 weights
diffusion_state = paddle.load(str(ckpt_dir / "best_diffusion_snapshot.pdparams"))
model.atom_coord_diffusion_model.set_state_dict(diffusion_state)
print(f"Loaded diffusion weights: {len(diffusion_state)} keys")
lattice_state = paddle.load(str(ckpt_dir / "best_lattice_snapshot.pdparams"))
model.lattice_sampler.set_state_dict(lattice_state)
print(f"Loaded lattice weights: {len(lattice_state)} keys")
sg_state = paddle.load(str(ckpt_dir / "best_space_group_snapshot.pdparams"))
model.space_group_sampler.set_state_dict(sg_state)
print(f"Loaded space_group weights: {len(sg_state)} keys")
wyckoff_state = paddle.load(str(ckpt_dir / "best_wyckoff-transformer_snapshot.pdparams"))
model.wyckoff_and_element_sampler.set_state_dict(wyckoff_state)
print(f"Loaded wyckoff weights: {len(wyckoff_state)} keys")
# Prepare inputs
lattice_lengths = paddle.to_tensor(FIXED_SAMPLE["lattice_lengths"], dtype="float32").unsqueeze(0)
lattice_angles = paddle.to_tensor(FIXED_SAMPLE["lattice_angles"], dtype="float32").unsqueeze(0)
wyckoff_indices = paddle.to_tensor(FIXED_SAMPLE["wyckoff_indices"], dtype="int64")
element_indices = paddle.to_tensor(FIXED_SAMPLE["element_indices"], dtype="int64")
space_group_indices = paddle.to_tensor([FIXED_SAMPLE["space_group_number"] - 1], dtype="int64")
n_asu = paddle.to_tensor([len(FIXED_SAMPLE["wyckoff_indices"])], dtype="int64")
frac_coords = paddle.to_tensor(FIXED_SAMPLE["frac_coords"], dtype="float32")
lattice_matrices = lattice_params_to_matrix_paddle(lattice_lengths, lattice_angles)
results = {}
# ---- Test 1: Diffusion model ----
diffusion_results = {}
dm = model.atom_coord_diffusion_model
fort_valin TIMESTEPS:
t_tensor = paddle.to_tensor([float(t_val)], dtype="float32")
time_embeddings = dm.time_embedder(t_tensor).expand([4, -1])
pred = dm.predict_equivariant_vectors(
time_embeddings=time_embeddings,
frac_coords=frac_coords,
element_indices=element_indices,
wyckoff_indices=wyckoff_indices,
space_group_indices=space_group_indices,
n_atoms_per_xtal=n_asu,
lattice_matrices=lattice_matrices,
lattice_lengths=lattice_lengths,
lattice_angles=lattice_angles,
)
diffusion_results[str(t_val)] = pred.numpy().tolist()
print(f"[diffusion] t={t_val}: pred abs max = {float(paddle.abs(pred).max().numpy()):.8f}")
results["diffusion"] = diffusion_results
# ---- Test 2: SpaceGroupSampler ----
sg_log_prob = model.space_group_sampler.log_prob(space_group_indices)
results["space_group"] = {
"log_prob": sg_log_prob.numpy().tolist(),
}
print(f"[space_group] log_prob = {float(sg_log_prob.numpy()):.8f}")
# ---- Test 3: LatticeSampler ----
paddle.seed(42)
lengths, angles, log_pfs, regularizer = model.lattice_sampler(space_group_indices)
results["lattice"] = {
"lengths": lengths.numpy().tolist(),
"angles": angles.numpy().tolist(),
"log_pfs": log_pfs.numpy().tolist(),
"regularizer": regularizer.numpy().tolist(),
}
print(f"[lattice] lengths={lengths.numpy().tolist()}, angles={angles.numpy().tolist()}")
print(f"[lattice] log_pfs={float(log_pfs.numpy()):.8f}, regularizer={float(regularizer.numpy()):.8f}")
# ---- Test 4: WyckoffElementTransformer ----
paddle.seed(42)
(
we_element_indices,
we_wyckoff_indices,
we_n_asu,
elements_log_prob,
wyckoffs_log_prob,
termination_log_prob,
) = model.wyckoff_and_element_sampler.sample_and_log_prob(
lattice_lengths=lattice_lengths,
lattice_angles=lattice_angles,
space_group_indices=space_group_indices,
temperature=1.0,
)
results["wyckoff"] = {
"element_indices": we_element_indices.numpy().tolist(),
"wyckoff_indices": we_wyckoff_indices.numpy().tolist(),
"n_asu_atoms_per_xtal": we_n_asu.numpy().tolist(),
"elements_log_prob": elements_log_prob.numpy().tolist(),
"wyckoffs_log_prob": wyckoffs_log_prob.numpy().tolist(),
"termination_log_prob": termination_log_prob.numpy().tolist(),
}
print(f"[wyckoff] n_atoms={we_n_asu.numpy().tolist()}")
print(f"[wyckoff] element_indices={we_element_indices.numpy().tolist()}")
print(f"[wyckoff] wyckoff_indices={we_wyckoff_indices.numpy().tolist()}")
output = {"framework": "pd", "results": results}
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w") as f:
json.dump(output, f, indent=2)
print(f"Saved to {output_path}")
def run_pt(args):
import os
import numpy as np
import torch
from utils.io_utils import convert_wandb_config_to_hydra_config
from utils.train_utils import dispatch_model, experiment_setup
from utils.data_utils import lattice_params_to_matrix_torch
torch.manual_seed(42)
np.random.seed(42)
PROJECT_ROOT = Path(os.getcwd())
ckpt_dir = str(PROJECT_ROOT / "pretrained-pt/mp_20")
config = convert_wandb_config_to_hydra_config(os.path.join(ckpt_dir, "config.yaml"))
experiment_setup(config)
device = torch.device("cpu")
model = dispatch_model(config, device)
forsnapshot_pathin sorted(Path(ckpt_dir).glob("best*snapshot*")):
p = snapshot_path.as_posix()
sd = torch.load(p, map_location=device, weights_only=True)
if"best_space_group_snapshot.pt"in p:
model.space_group_sampler.load_state_dict(sd)
elif"best_lattice_snapshot.pt"in p:
model.lattice_sampler.load_state_dict(sd)
elif"best_wyckoff-transformer_snapshot.pt"in p:
model.wyckoff_and_element_sampler.load_state_dict(sd)
elif"best_diffusion_snapshot.pt"in p:
model.atom_coord_diffusion_model.load_state_dict(sd)
model.eval()
lattice_lengths = torch.tensor(FIXED_SAMPLE["lattice_lengths"], dtype=torch.float32, device=device).unsqueeze(0)
lattice_angles = torch.tensor(FIXED_SAMPLE["lattice_angles"], dtype=torch.float32, device=device).unsqueeze(0)
wyckoff_indices = torch.tensor(FIXED_SAMPLE["wyckoff_indices"], dtype=torch.long, device=device)
element_indices = torch.tensor(FIXED_SAMPLE["element_indices"], dtype=torch.long, device=device)
space_group_indices = torch.tensor([FIXED_SAMPLE["space_group_number"] - 1], dtype=torch.long, device=device)
n_asu = torch.tensor([len(FIXED_SAMPLE["wyckoff_indices"])], dtype=torch.long, device=device)
frac_coords = torch.tensor(FIXED_SAMPLE["frac_coords"], dtype=torch.float32, device=device)
lattice_matrices = lattice_params_to_matrix_torch(lattice_lengths, lattice_angles)
results = {}
# ---- Test 1: Diffusion model ----
dm = model.atom_coord_diffusion_model
dm.eval()
diffusion_results = {}
fort_valin TIMESTEPS:
t_tensor = torch.tensor([float(t_val)], dtype=torch.float32, device=device)
time_embeddings = dm.time_embedder(t_tensor).expand([4, -1])
pred = dm.predict_equivariant_vectors(
time_embeddings=time_embeddings,
frac_coords=frac_coords,
element_indices=element_indices,
wyckoff_indices=wyckoff_indices,
space_group_indices=space_group_indices,
n_atoms_per_xtal=n_asu,
lattice_matrices=lattice_matrices,
lattice_lengths=lattice_lengths,
lattice_angles=lattice_angles,
)
diffusion_results[str(t_val)] = pred.detach().cpu().numpy().tolist()
print(f"[diffusion] t={t_val}: pred abs max = {float(pred.abs().max()):.8f}")
results["diffusion"] = diffusion_results
# ---- Test 2: SpaceGroupSampler ----
sg_log_prob = model.space_group_sampler.log_prob(space_group_indices)
results["space_group"] = {
"log_prob": sg_log_prob.detach().cpu().numpy().tolist(),
}
print(f"[space_group] log_prob = {float(sg_log_prob):.8f}")
# ---- Test 3: LatticeSampler ----
torch.manual_seed(42)
lengths, angles, log_pfs, regularizer = model.lattice_sampler(space_group_indices)
results["lattice"] = {
"lengths": lengths.detach().cpu().numpy().tolist(),
"angles": angles.detach().cpu().numpy().tolist(),
"log_pfs": log_pfs.detach().cpu().numpy().tolist(),
"regularizer": regularizer.detach().cpu().numpy().tolist(),
}
print(f"[lattice] lengths={lengths.detach().cpu().numpy().tolist()}, angles={angles.detach().cpu().numpy().tolist()}")
print(f"[lattice] log_pfs={float(log_pfs):.8f}, regularizer={float(regularizer):.8f}")
# ---- Test 4: WyckoffElementTransformer ----
torch.manual_seed(42)
(
we_element_indices,
we_wyckoff_indices,
we_n_asu,
elements_log_prob,
wyckoffs_log_prob,
termination_log_prob,
) = model.wyckoff_and_element_sampler.sample_and_log_prob(
lattice_lengths=lattice_lengths,
lattice_angles=lattice_angles,
space_group_indices=space_group_indices,
temperature=1.0,
)
results["wyckoff"] = {
"element_indices": we_element_indices.detach().cpu().numpy().tolist(),
"wyckoff_indices": we_wyckoff_indices.detach().cpu().numpy().tolist(),
"n_asu_atoms_per_xtal": we_n_asu.detach().cpu().numpy().tolist(),
"elements_log_prob": elements_log_prob.detach().cpu().numpy().tolist(),
"wyckoffs_log_prob": wyckoffs_log_prob.detach().cpu().numpy().tolist(),
"termination_log_prob": termination_log_prob.detach().cpu().numpy().tolist(),
}
print(f"[wyckoff] n_atoms={we_n_asu.detach().cpu().numpy().tolist()}")
print(f"[wyckoff] element_indices={we_element_indices.detach().cpu().numpy().tolist()}")
print(f"[wyckoff] wyckoff_indices={we_wyckoff_indices.detach().cpu().numpy().tolist()}")
output = {"framework": "pt", "results": results}
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w") as f:
json.dump(output, f, indent=2)
print(f"Saved to {output_path}")
STOCHASTIC_KEYS = {
"lattice.lengths",
"lattice.log_pfs",
"wyckoff.element_indices",
"wyckoff.wyckoff_indices",
"wyckoff.n_asu_atoms_per_xtal",
"wyckoff.elements_log_prob",
"wyckoff.wyckoffs_log_prob",
"wyckoff.termination_log_prob",
}
def _compare_dict(pt_dict, pd_dict, path=""):
""" Recursively compare two dicts containing numeric arrays. Returns list of (path, max_diff, mean_diff, status, is_stochastic) tuples."""
import numpy as np
rows = []
forkeyinsorted(pt_dict.keys()):
pt_val = pt_dict[key]
pd_val = pd_dict.get(key)
if pd_val is None:
rows.append((f"{path}.{key}", None, None, "MISSING_IN_PD", False))continue
current_path = f"{path}.{key}"if path else key
if isinstance(pt_val, dict):
rows.extend(_compare_dict(pt_val, pd_val, current_path))
else:
pt_arr = np.array(pt_val)
pd_arr = np.array(pd_val)
is_stochastic = current_path in STOCHASTIC_KEYS
if pt_arr.shape != pd_arr.shape:
status = "RNG_DIFF (shape mismatch, expected)"if is_stochastic else f"SHAPE_MISMATCH: PT{pt_arr.shape} vs PD{pd_arr.shape}"
rows.append((current_path, None, None, status, is_stochastic))continue
diff = np.abs(pt_arr - pd_arr)
max_diff = float(diff.max())
mean_diff = float(diff.mean())
if is_stochastic:
status = "RNG_DIFF (expected due to different Categorical/multinomial RNG)"
else:
status = "PASS"if max_diff < 1e-4 else"FAIL"
rows.append((current_path, max_diff, mean_diff, status, is_stochastic))return rows
def run_diff(args):
import numpy as np
with open(args.pt) as f:
pt = json.load(f)
with open(args.pd) as f:
pd = json.load(f)
print("="* 70)
print("Forward Alignment Diff Summary")
print("="* 70)
pt_results = pt.get("results", {})
pd_results = pd.get("results", {})
all_rows = _compare_dict(pt_results, pd_results)
# Print diffusion results first
diffusion_rows = [r forrin all_rows if r[0].startswith("diffusion.")]
other_rows = [r forrin all_rows if not r[0].startswith("diffusion.")]
if diffusion_rows:
print("\n--- Diffusion Model (predict_equivariant_vectors) ---")
forpath, max_diff, mean_diff, status, is_stochin diffusion_rows:
t_val = path.split(".")[-1]
print(f" t={t_val}: max|diff|={max_diff:.2e} mean|diff|={mean_diff:.2e} [{status}]")
formodule_namein ["space_group", "lattice", "wyckoff"]:
module_rows = [r forrin other_rows if r[0].startswith(f"{module_name}.")]
if module_rows:
print(f"\n--- {module_name.replace('_', ' ').title()} ---")
forpath, max_diff, mean_diff, status, is_stochin module_rows:
key = path.split(".", 1)[1]
if max_diff is None:
print(f" {key}: {status}")
else:
tag = " (stochastic)"if is_stoch else""
print(f" {key}: max|diff|={max_diff:.2e} mean|diff|={mean_diff:.2e}{tag}")
print(f" [{status}]")
print("\n" + "="* 70)
# Only deterministic outputs count toward overall PASS/FAIL
deterministic_rows = [r forrin all_rows if not r[4] and r[1] is not None]
stochastic_rows = [r forrin all_rows if r[4]]
if deterministic_rows:
det_max_diffs = [r[1] forrin deterministic_rows]
overall_max = max(det_max_diffs)
print(f"Deterministic outputs max |diff|: {overall_max:.2e}")
if overall_max < 1e-4:
print("PASS: All deterministic modules within 1e-4 threshold")
else:
print("FAIL: Some deterministic modules exceed 1e-4 threshold")
else:
print("No deterministic numeric values found.")
if stochastic_rows:
print()
print("Note: Stochastic outputs (lattice.lengths, lattice.log_pfs, wyckoff.*)")
print(" differ due to PyTorch vs PaddlePaddle Categorical/multinomial RNG.")
print(" This is expected and not a precision bug.")
def main():
parser = argparse.ArgumentParser()
sub = parser.add_subparsers(dest="mode")
pd_p = sub.add_parser("pd")
pd_p.add_argument("--output", required=True)
pt_p = sub.add_parser("pt")
pt_p.add_argument("--output", required=True)
diff_p = sub.add_parser("diff")
diff_p.add_argument("--pt", required=True)
diff_p.add_argument("--pd", required=True)
args = parser.parse_args()
if args.mode == "pd":
run_pd(args)
elif args.mode == "pt":
run_pt(args)
elif args.mode == "diff":
run_diff(args)
else:
parser.print_help()
if __name__ == "__main__":
main()3.2 反向精度对齐使用固定 crystal 输入 + 固定 target score 进行 AdamW 微训练，对比 PT/PD loss 轨迹和梯度范数趋势。
3.2.1 反向对齐loss对比
反向对齐loss的Python脚本,点击后展开脚本#!/usr/bin/env python3# -*- coding: utf-8 -*-"""反向对齐（微训练）脚本：训练 >=2 轮，比较 PT/PD loss 轨迹。说明：- 使用固定 crystal 输入 + 固定 target score，进行可复现的 AdamW 训练。- 该脚本聚焦“反向与优化器路径是否一致”，避免数据管线/RNG 差异干扰。"""
import argparse
import json
import math
import sys
from pathlib import Path
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[5]
FIXED_SAMPLE = {
"space_group_number": 139,
"lattice_lengths": [5.141545295715332, 5.141545295715332, 9.469661712646484],
"lattice_angles": [90.0, 90.0, 90.0],
"element_indices": [64, 26, 13, 13], # Dy, Co, Si, Si (0-indexed)"wyckoff_indices": [0, 5, 8, 9],
"frac_coords": [
[0.0000, 0.0000, 0.0000],
[0.2500, 0.2500, 0.2500],
[0.0000, 0.3354, 0.0000],
[0.2688, 0.5000, 0.0000],
],
}
FIXED_TARGET_SCORE = [
[0.0100, -0.0200, 0.0150],
[-0.0120, 0.0080, -0.0060],
[0.0070, 0.0110, -0.0130],
[-0.0040, -0.0090, 0.0120],
]
def _resolve_path(p: str) -> Path:
path = Path(p)
if not path.is_absolute():
path = PROJECT_ROOT / path
return path
def _build_t_schedule(steps_per_epoch: int, fixed_timestep: int):
return [int(fixed_timestep) for_in range(steps_per_epoch)]
def _save_json(save_path: Path, data: dict):
save_path.parent.mkdir(parents=True, exist_ok=True)
with open(save_path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2, ensure_ascii=False)
def _select_named_params(named_params, keyword: str):
selected = [(n, p) forn, pin named_params if keyword in n]
if not selected:
raise ValueError(f"No parameters matched keyword: {keyword}")
return selected
def run_pt(args):
import torch
import torch.nn.functional as F
from utils.train_utils import dispatch_model, experiment_setup
from utils.io_utils import convert_wandb_config_to_hydra_config
from utils.data_utils import lattice_params_to_matrix_torch
torch.manual_seed(args.seed)
np.random.seed(args.seed)
ckpt_dir = _resolve_path(args.ckpt_dir)
config = convert_wandb_config_to_hydra_config(str(ckpt_dir / "config.yaml"))
experiment_setup(config)
device = torch.device("cpu")
model = dispatch_model(config, device)
forsnapshot_pathin sorted(ckpt_dir.glob("best*snapshot*")):
p = snapshot_path.as_posix()
sd = torch.load(p, map_location=device, weights_only=True)
if"best_space_group_snapshot.pt"in p:
model.space_group_sampler.load_state_dict(sd)
elif"best_lattice_snapshot.pt"in p:
model.lattice_sampler.load_state_dict(sd)
elif"best_wyckoff-transformer_snapshot.pt"in p:
model.wyckoff_and_element_sampler.load_state_dict(sd)
elif"best_diffusion_snapshot.pt"in p:
model.atom_coord_diffusion_model.load_state_dict(sd)
dm = model.atom_coord_diffusion_model
dm.eval()
lattice_lengths = torch.tensor(FIXED_SAMPLE["lattice_lengths"], dtype=torch.float32, device=device).unsqueeze(0)
lattice_angles = torch.tensor(FIXED_SAMPLE["lattice_angles"], dtype=torch.float32, device=device).unsqueeze(0)
wyckoff_indices = torch.tensor(FIXED_SAMPLE["wyckoff_indices"], dtype=torch.long, device=device)
element_indices = torch.tensor(FIXED_SAMPLE["element_indices"], dtype=torch.long, device=device)
space_group_indices = torch.tensor([FIXED_SAMPLE["space_group_number"] - 1], dtype=torch.long, device=device)
n_asu = torch.tensor([len(FIXED_SAMPLE["wyckoff_indices"])], dtype=torch.long, device=device)
frac_coords = torch.tensor(FIXED_SAMPLE["frac_coords"], dtype=torch.float32, device=device)
target_score = torch.tensor(FIXED_TARGET_SCORE, dtype=torch.float32, device=device)
lattice_matrices = lattice_params_to_matrix_torch(lattice_lengths, lattice_angles)
selected_named = _select_named_params(list(dm.named_parameters()), args.train_param_keyword)
selected_params = [p for_, pin selected_named]
optimizer = torch.optim.AdamW(
selected_params,
lr=args.lr,
betas=(args.beta1, args.beta2),
eps=args.eps,
weight_decay=args.weight_decay,
)
t_schedule = _build_t_schedule(args.steps_per_epoch, args.fixed_timestep)
records = []
global_step = 0
forepochin range(1, args.epochs + 1):
forstep_in_epoch, t_valin enumerate(t_schedule, start=1):
t_tensor = torch.tensor([float(t_val)], dtype=torch.float32, device=device)
n_asu_atoms = frac_coords.shape[0]
time_embeddings = dm.time_embedder(t_tensor).expand([n_asu_atoms, -1])
pred = dm.predict_equivariant_vectors(
time_embeddings=time_embeddings,
frac_coords=frac_coords,
element_indices=element_indices,
wyckoff_indices=wyckoff_indices,
space_group_indices=space_group_indices,
n_atoms_per_xtal=n_asu,
lattice_matrices=lattice_matrices,
lattice_lengths=lattice_lengths,
lattice_angles=lattice_angles,
)
loss = F.mse_loss(pred, target_score)
if not torch.isfinite(loss):
raise RuntimeError(
f"PT loss is non-finite at epoch={epoch}, step={step_in_epoch}, t={t_val}"
)
optimizer.zero_grad(set_to_none=True)
loss.backward()
grad_sq = 0.0
forpindm.parameters():
if p.grad is not None:
grad_sq += float((p.grad.detach() **2).sum().item())
grad_norm = math.sqrt(grad_sq)
optimizer.step()
rec = {
"global_step": global_step,
"epoch": epoch,
"step_in_epoch": step_in_epoch,
"timestep": int(t_val),
"loss": float(loss.detach().item()),
"grad_norm": float(grad_norm),
}
records.append(rec)
print(
f"[PT] epoch={epoch} step={step_in_epoch}/{args.steps_per_epoch} "
f"global={global_step} t={t_val} loss={rec['loss']:.8f} grad={grad_norm:.8f}"
)
global_step += 1
out = {
"framework": "pt",
"seed": args.seed,
"epochs": args.epochs,
"steps_per_epoch": args.steps_per_epoch,
"optimizer": {
"name": "AdamW",
"lr": args.lr,
"beta1": args.beta1,
"beta2": args.beta2,
"eps": args.eps,
"weight_decay": args.weight_decay,
},
"train_param_keyword": args.train_param_keyword,
"trained_params": [n forn, _in selected_named],
"records": records,
}
save_path = _resolve_path(args.save_json)
_save_json(save_path, out)
print(f"Saved PT training trace: {save_path}")
def run_pd(args):
import paddle
import paddle.nn.functional as F
sys.path.insert(0, str(PROJECT_ROOT))
from omegaconf import OmegaConf
from ppmat.models.sgequidiff.diffusion_model import (
EquivariantDiffusionModel,
EquivariantDiffusionModelConfig,
)
from ppmat.models.sgequidiff.non_equivariant_drift_modules import GNNConfig
from ppmat.models.sgequidiff.embedding_utils import set_global_embedding_tools
from ppmat.models.sgequidiff.data_utils import lattice_params_to_matrix_paddle
paddle.seed(args.seed)
np.random.seed(args.seed)
ckpt_dir = _resolve_path(args.ckpt_dir)
set_global_embedding_tools(element_embedding_json_path="cgcnn_atom_init.json")
gnn_config = GNNConfig()
gnn_config.dataset_name = "mp_20"
gnn_config.num_plane_wave_freqs = 96
gnn_config.num_cartesian_distance_gaussians = 96
gnn_config.atom_hidden_dim = 256
gnn_config.edge_hidden_dim = 128
gnn_config.num_msg_pass_steps = 5
gnn_config.use_vpa = True
gnn_config.use_graph_norm = True
gnn_config.use_frac_coords_in_node_emb = True
gnn_config.cutoff = 10.0
cfg = EquivariantDiffusionModelConfig(
model_type="gnn",
num_wn_lattice_translations=3,
num_timesteps=1000,
noise_scheduler_num_monte_carlo_samples=2500,
time_emb_dim=128,
gnn_config=gnn_config,
)
model = EquivariantDiffusionModel(cfg)
model.eval()
dm = model # dm is the diffusion model we'll use for training
weight_files = [
("diffusion", "best_diffusion_snapshot.pdparams"),
("lattice", "best_lattice_snapshot.pdparams"),
("space_group", "best_space_group_snapshot.pdparams"),
("wyckoff", "best_wyckoff-transformer_snapshot.pdparams"),
]
forname, filenamein weight_files:
weight_path = ckpt_dir / filename
ifweight_path.exists():
print(f"Loading {name} weights from {filename}")
state = paddle.load(str(weight_path))
# Note: Weights are already correctly converted by convert_sgequidiff_weights.py# No additional transposition needed here# For diffusion model, load directly into modelif name == "diffusion":
model.set_state_dict(state)
# Skip other weights for this test
print(f" Loaded {len(state)} keys")
else:
print(f"Warning: {filename} not found in {ckpt_dir}")
dm.eval()
print("Model loaded successfully")
lattice_lengths = paddle.to_tensor(FIXED_SAMPLE["lattice_lengths"], dtype=paddle.float32).unsqueeze(0)
lattice_angles = paddle.to_tensor(FIXED_SAMPLE["lattice_angles"], dtype=paddle.float32).unsqueeze(0)
wyckoff_indices = paddle.to_tensor(FIXED_SAMPLE["wyckoff_indices"], dtype=paddle.int64)
element_indices = paddle.to_tensor(FIXED_SAMPLE["element_indices"], dtype=paddle.int64)
space_group_indices = paddle.to_tensor([FIXED_SAMPLE["space_group_number"] - 1], dtype=paddle.int64)
n_asu = paddle.to_tensor([len(FIXED_SAMPLE["wyckoff_indices"])], dtype=paddle.int64)
frac_coords = paddle.to_tensor(FIXED_SAMPLE["frac_coords"], dtype=paddle.float32)
target_score = paddle.to_tensor(FIXED_TARGET_SCORE, dtype=paddle.float32)
lattice_matrices = lattice_params_to_matrix_paddle(lattice_lengths, lattice_angles)
selected_named = _select_named_params(list(dm.named_parameters()), args.train_param_keyword)
selected_params = [p for_, pin selected_named]
optimizer = paddle.optimizer.AdamW(
learning_rate=args.lr,
beta1=args.beta1,
beta2=args.beta2,
epsilon=args.eps,
weight_decay=args.weight_decay,
parameters=selected_params,
)
t_schedule = _build_t_schedule(args.steps_per_epoch, args.fixed_timestep)
records = []
global_step = 0
forepochin range(1, args.epochs + 1):
forstep_in_epoch, t_valin enumerate(t_schedule, start=1):
t_tensor = paddle.to_tensor([float(t_val)], dtype=paddle.float32)
n_asu_atoms = frac_coords.shape[0]
time_embeddings = dm.time_embedder(t_tensor).expand([n_asu_atoms, -1])
pred = dm.predict_equivariant_vectors(
time_embeddings=time_embeddings,
frac_coords=frac_coords,
element_indices=element_indices,
wyckoff_indices=wyckoff_indices,
space_group_indices=space_group_indices,
n_atoms_per_xtal=n_asu,
lattice_matrices=lattice_matrices,
lattice_lengths=lattice_lengths,
lattice_angles=lattice_angles,
)
loss = F.mse_loss(pred, target_score)
if not bool(paddle.isfinite(loss).numpy()):
raise RuntimeError(
f"PD loss is non-finite at epoch={epoch}, step={step_in_epoch}, t={t_val}"
)
optimizer.clear_grad()
loss.backward(retain_graph=True)
grad_sq = 0.0
forpindm.parameters():
if p.grad is not None:
grad_val = p.grad.detach()
if not bool(paddle.isfinite(grad_val).all().numpy()):
continue
grad_sq += float((grad_val **2).sum().numpy())
grad_norm = math.sqrt(grad_sq)
optimizer.step()
rec = {
"global_step": global_step,
"epoch": epoch,
"step_in_epoch": step_in_epoch,
"timestep": int(t_val),
"loss": float(loss.detach().numpy()),
"grad_norm": float(grad_norm),
}
records.append(rec)
print(
f"[PD] epoch={epoch} step={step_in_epoch}/{args.steps_per_epoch} "
f"global={global_step} t={t_val} loss={rec['loss']:.8f} grad={grad_norm:.8f}"
)
global_step += 1
out = {
"framework": "pd",
"seed": args.seed,
"epochs": args.epochs,
"steps_per_epoch": args.steps_per_epoch,
"optimizer": {
"name": "AdamW",
"lr": args.lr,
"beta1": args.beta1,
"beta2": args.beta2,
"eps": args.eps,
"weight_decay": args.weight_decay,
},
"train_param_keyword": args.train_param_keyword,
"trained_params": [n forn, _in selected_named],
"records": records,
}
save_path = _resolve_path(args.save_json)
_save_json(save_path, out)
print(f"Saved PD training trace: {save_path}")
def run_diff(args):
pt_path = _resolve_path(args.pt_json)
pd_path = _resolve_path(args.pd_json)
with open(pt_path, "r", encoding="utf-8") as f:
pt = json.load(f)
with open(pd_path, "r", encoding="utf-8") as f:
pd = json.load(f)
pt_rec = pt["records"]
pd_rec = pd["records"]
if len(pt_rec) != len(pd_rec):
raise ValueError(f"records size mismatch: PT={len(pt_rec)} PD={len(pd_rec)}")
loss_diffs = []
grad_rel_diffs = []
rows = []
fora, bin zip(pt_rec, pd_rec):
loss_diff = abs(a["loss"] - b["loss"])
grad_rel = abs(a["grad_norm"] - b["grad_norm"]) / (abs(a["grad_norm"]) + 1e-12)
loss_diffs.append(loss_diff)
grad_rel_diffs.append(grad_rel)
rows.append(
{
"global_step": a["global_step"],
"epoch": a["epoch"],
"step_in_epoch": a["step_in_epoch"],
"timestep": a["timestep"],
"pt_loss": a["loss"],
"pd_loss": b["loss"],
"abs_loss_diff": loss_diff,
"pt_grad_norm": a["grad_norm"],
"pd_grad_norm": b["grad_norm"],
"rel_grad_norm_diff": grad_rel,
}
)
max_loss_diff = float(np.max(loss_diffs)) if loss_diffs else 0.0
mean_loss_diff = float(np.mean(loss_diffs)) if loss_diffs else 0.0
max_rel_grad_diff = float(np.max(grad_rel_diffs)) if grad_rel_diffs else 0.0
print("="* 80)
print("Backward Align Diff Summary")
print("="* 80)
print(f"PT records: {len(pt_rec)} | PD records: {len(pd_rec)}")
print(f"max |loss diff| : {max_loss_diff:.8e}")
print(f"mean|loss diff| : {mean_loss_diff:.8e}")
print(f"max rel grad diff : {max_rel_grad_diff:.4%}")
threshold = args.loss_threshold
if max_loss_diff < threshold:
print(f"PASS: max |loss diff| < {threshold}")
else:
print(f"FAIL: max |loss diff| >= {threshold}")
out = {
"pt_json": str(pt_path),
"pd_json": str(pd_path),
"num_steps": len(rows),
"max_abs_loss_diff": max_loss_diff,
"mean_abs_loss_diff": mean_loss_diff,
"max_rel_grad_norm_diff": max_rel_grad_diff,
"loss_threshold": threshold,
"pass": bool(max_loss_diff < threshold),
"rows": rows,
}
save_path = _resolve_path(args.save_json)
_save_json(save_path, out)
print(f"Saved diff summary: {save_path}")
def main():
parser = argparse.ArgumentParser()
sub = parser.add_subparsers(dest="mode")
def add_common_train_args(p):
p.add_argument("--ckpt_dir", required=True)
p.add_argument("--save_json", required=True)
p.add_argument("--epochs", type=int, default=2)
p.add_argument("--steps_per_epoch", type=int, default=4)
p.add_argument("--fixed_timestep", type=int, default=100)
p.add_argument("--train_param_keyword", type=str, default="non_equivariant_drift_model.mlp_out")
p.add_argument("--seed", type=int, default=42)
p.add_argument("--lr", type=float, default=1e-3)
p.add_argument("--beta1", type=float, default=0.9)
p.add_argument("--beta2", type=float, default=0.999)
p.add_argument("--eps", type=float, default=1e-8)
p.add_argument("--weight_decay", type=float, default=0.0)
pt_p = sub.add_parser("pt")
add_common_train_args(pt_p)
pd_p = sub.add_parser("pd")
add_common_train_args(pd_p)
diff_p = sub.add_parser("diff")
diff_p.add_argument("--pt_json", required=True)
diff_p.add_argument("--pd_json", required=True)
diff_p.add_argument("--save_json", required=True)
diff_p.add_argument("--loss_threshold", type=float, default=1e-3)
args = parser.parse_args()
if args.mode == "pt":
run_pt(args)
elif args.mode == "pd":
run_pd(args)
elif args.mode == "diff":
run_diff(args)
else:
parser.print_help()
if __name__ == "__main__":
main()
3.3 采样指标对齐采样指标对齐，在下一个评论里；github禁止PR超过64K，拆分2段！！！ |
learncat163
commented
Jun 24, 2026
Author
3.3 采样指标对齐采样指标对接结果
采样的对齐步骤比较复杂，主要分离为2步：
3.3.1 PT的晶体采样PT晶体采样Python脚本,点击后展开脚本#!/usr/bin/env python3# -*- coding: utf-8 -*-"""PyTorch 原始环境晶体生成脚本。在原始 sgequidiff 的venv 环境中运行，使用原始 PyTorch 代码生成晶体样本。使用方法: source ~/opensource/cailiao/sgequidiff/.venv/bin/activate python generate_samples_pt.py --num_samples 64 --output_dir ../outputs/pt_samples"""importargparseimportjsonimportosimportsysimporttimefrompathlibimportPathdefparse_args():
parser=argparse.ArgumentParser(description="PT 晶体生成")
parser.add_argument("--num_samples", type=int, default=64, help="生成样本数")
parser.add_argument("--batch_size", type=int, default=32, help="批大小")
parser.add_argument("--temperature", type=float, default=1.0, help="采样温度")
parser.add_argument("--seed", type=int, default=42, help="随机种子")
parser.add_argument("--ckpt_dir", type=str, required=True, help="PT 权重目录")
parser.add_argument("--output_dir", type=str, required=True, help="输出目录")
parser.add_argument("--dataset", type=str, default="mp_20", choices=["mp_20", "mpts_52"])
returnparser.parse_args()
defmain():
args=parse_args()
output_dir=Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# 记录日志log_path=output_dir/"generation.log"sys.stdout=open(log_path, "w")
sys.stderr=sys.stdoutprint(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] PT 晶体生成开始")
print(f" 参数: num_samples={args.num_samples}, batch_size={args.batch_size}, "f"temperature={args.temperature}, seed={args.seed}")
print(f" 权重: {args.ckpt_dir}")
print(f" 输出: {args.output_dir}")
# 添加原始项目路径（从脚本位置上溯到项目根）raw_project=Path(__file__).resolve().parents[4] /"sgequidiff-raw"sys.path.insert(0, str(raw_project))
os.chdir(str(raw_project))
# 导入原始模块fromscripts.generate_crystalsimportmainasraw_main# 构造 raw_main 的参数sys.argv= [
"generate_crystals.py",
"--num_samples", str(args.num_samples),
"--batch_size", str(args.batch_size),
"--ckpt_dir", args.ckpt_dir,
"--load_best_submodules",
"--temperature", str(args.temperature),
"--seed", str(args.seed),
"--save_method", "cif",
"--save_dir", str(output_dir/"cifs"),
]
ifargs.dataset=="mpts_52":
sys.argv+= ["--dataset_name", "mpts_52"]
print(f"\n启动原始生成脚本...")
raw_main()
# 保存元数据meta= {
"framework": "pt",
"dataset": args.dataset,
"num_samples": args.num_samples,
"temperature": args.temperature,
"seed": args.seed,
"ckpt_dir": args.ckpt_dir,
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
}
withopen(output_dir/"metadata.json", "w") asf:
json.dump(meta, f, indent=2)
print(f"\n[{time.strftime('%Y-%m-%d %H:%M:%S')}] PT 晶体生成完成")
print(f" 输出目录: {args.output_dir}")
if__name__=="__main__":
main()3.3.2 PD的晶体采样PD晶体采样Python脚本,点击后展开脚本#!/usr/bin/env python3# -*- coding: utf-8 -*-"""Paddle 环境晶体生成脚本（独立进程隔离版）。每生成一个晶体，都启动一个独立子进程，进程退出后操作系统自动回收全部 GPU 显存。 不然会爆显存 OOM！！！！主进程本身不加载模型、不分配 GPU 资源，仅负责调度和日志记录。架构: main() -- 主进程: 调度循环，不碰 GPU | +-- subprocess python generate_samples_pd.py --worker 0 +-- subprocess python generate_samples_pd.py --worker 1 +-- ... 每个子进程: 加载模型 -> 采样 -> 保存 CIF -> 进程退出 -> GPU 显存彻底释放使用方法: conda activate ppmat python generate_samples_pd.py --num_samples 64 --output_dir ./output"""importargparseimportjsonimportosimportsubprocessimportsysimporttimefrompathlibimportPathdefparse_args():
parser=argparse.ArgumentParser(
description="PD 晶体生成（独立进程隔离，每晶体独立子进程）")
parser.add_argument("--num_samples", type=int, default=16,
help="生成样本数 (默认 16)")
parser.add_argument("--temperature", type=float, default=1.0,
help="采样温度 (默认 1.0)")
parser.add_argument("--seed", type=int, default=42,
help="随机种子 (默认 42)")
parser.add_argument("--output_dir", type=str, required=True,
help="输出目录")
parser.add_argument("--dataset", type=str, default="mp_20",
choices=["mp_20", "mpts_52"],
help="数据集 (默认 mp_20)")
parser.add_argument("--diffusion_snr", type=float, default=0.4,
help="扩散 SNR (默认 0.4)")
parser.add_argument("--max_step_size", type=float, default=1e6,
help="扩散采样最大步长")
# --- worker 模式（由主进程自动调用，用户无需手动指定） ---parser.add_argument("--worker", type=int, default=None,
help="[内部] worker 模式: 生成第 N 个晶体后退出")
returnparser.parse_args()
# ---------------------------------------------------------------------------# Worker: 在独立子进程中运行，生成单个晶体后退出，GPU 显存随之彻底释放# ---------------------------------------------------------------------------defworker_main(args):
"""子进程入口: 构建模型 -> 采样 -> 保存 CIF -> 退出。"""project_root=str(Path(__file__).resolve().parents[4])
ifproject_rootnotinsys.path:
sys.path.insert(0, project_root)
importpaddlefromppmat.models.sgequidiff.weight_utilsimportload_pretrained_weightsfromppmat.models.sgequidiff.crystal_samplerimport (
CrystalSampler, CrystalSamplerConfig,
)
fromppmat.models.sgequidiff.diffusion_modelimport (
EquivariantDiffusionModelConfig,
)
fromppmat.models.sgequidiff.lattice_samplerimportLatticeSamplerConfigfromppmat.models.sgequidiff.wyckoff_transformerimport (
WyckoffElementTransformerConfig,
)
fromppmat.models.sgequidiff.non_equivariant_drift_modulesimportGNNConfigfromppmat.models.sgequidiff.embedding_utilsimport (
set_global_embedding_tools,
)
fromppmat.models.sgequidiff.constantsimport (
lattice_parameter_ranges, chemical_symbols,
)
idx=args.worker# 设置随机种子（每个 worker 不同）paddle.seed(args.seed+idx)
# 初始化全局 embedding 工具set_global_embedding_tools()
# 构建模型配置lr=lattice_parameter_ranges.get(
args.dataset, lattice_parameter_ranges["mp_20"],
)
gnn_cfg=GNNConfig(
num_plane_wave_freqs=96,
num_cartesian_distance_gaussians=96,
edge_hidden_dim=128,
atom_hidden_dim=256,
use_vpa=True,
use_graph_norm=True,
num_msg_pass_steps=5,
cutoff=10.0,
use_frac_coords_in_node_emb=True,
dataset_name=args.dataset,
)
diff_cfg=EquivariantDiffusionModelConfig(
model_type="gnn",
num_timesteps=1000,
noise_scheduler_num_monte_carlo_samples=2500,
num_wn_lattice_translations=3,
sigma_min=0.002,
sigma_max=0.5,
time_emb_dim=128,
num_plane_wave_freqs=96,
gnn_config=gnn_cfg,
)
lattice_cfg=LatticeSamplerConfig(
input_dimension=128,
hidden_dimension=256,
min_lattice_length=lr["min_lattice_length"],
max_lattice_length=lr["max_lattice_length"],
min_lattice_angle=lr["min_lattice_angle"],
max_lattice_angle=lr["max_lattice_angle"],
)
we_cfg=WyckoffElementTransformerConfig(
hidden_dim=256,
dataset_name=args.dataset,
num_heads=4,
num_hidden_layers=1,
dropout_rate=0.0,
)
sampler_cfg=CrystalSamplerConfig(
diffusion_model_config=diff_cfg,
lattice_model_config=lattice_cfg,
transformer_config=we_cfg,
)
# 构建模型 & 加载权重model=CrystalSampler(sampler_cfg)
model.eval()
load_pretrained_weights(model, dataset_name=args.dataset, verbose=False)
# 采样try:
crystals=model.sample_crystal(
batch_size=1,
diffusion_snr=args.diffusion_snr,
temperature=args.temperature,
)
exceptExceptionase:
print(f"FAIL:{idx}:sample_error:{e}", flush=True)
returnifnotcrystalsorlen(crystals) ==0:
print(f"FAIL:{idx}:empty_result", flush=True)
return# 保存 CIFc=crystals[0]
lattice= (
c.conventional_lattice_lengths.tolist()
+c.conventional_lattice_angles.tolist()
)
elements= [chemical_symbols[i] foriinc.element_indices.tolist()]
coords=c.conventional_frac_coords.tolist()
cif_path=Path(args.output_dir) /"cifs"/f"gen_{idx:05d}.cif"_save_cif(cif_path, lattice, elements, coords, c.space_group_number)
# 输出统计信息（主进程可捕获）real_atoms=sum(1foreinelementsifenotin ("X", "X0+", ""))
print(f"OK:{idx}:atoms={real_atoms}", flush=True)
# 显式清理（虽然进程即将退出，但保险起见）delmodel, crystalspaddle.device.cuda.empty_cache()
# ---------------------------------------------------------------------------# CIF 保存（混合占位模式，与 PT 版本对齐）# ---------------------------------------------------------------------------def_save_cif(path, lattice_params, elements, frac_coords, space_group):
"""将晶体信息保存为 CIF 文件。 相同分数坐标的原子合并为混合占位 (mixed-occupancy) 位点， 使 CIF 的位点数与 PT 版本一致，消除评估偏差。 """fromcollectionsimportCounter, defaultdictfrompymatgen.coreimportElement, Lattice, Structure# 过滤占位元素valid_pairs= []
fori, el_nameinenumerate(elements):
ifel_namein ("X", "X0+", "") orlen(el_name) ==0:
continuetry:
el=Element(el_name)
coord=tuple(frac_coords[i])
valid_pairs.append((el, coord))
exceptException:
continueifnotvalid_pairs:
return# 按坐标分组，合并相同坐标的不同元素为混合占位site_map=defaultdict(list)
forel, coordinvalid_pairs:
site_map[coord].append(el)
group_species= []
group_coords= []
forcoord, elements_listinsite_map.items():
el_counts=Counter(elements_list)
total=sum(el_counts.values())
iftotal==1:
group_species.append(elements_list[0])
else:
group_species.append(
{el: count/totalforel, countinel_counts.items()}
)
group_coords.append(list(coord))
a, b, c=lattice_params[0], lattice_params[1], lattice_params[2]
alpha, beta, gamma= (
lattice_params[3], lattice_params[4], lattice_params[5],
)
lattice=Lattice.from_parameters(a, b, c, alpha, beta, gamma)
structure=Structure(
lattice, group_species, group_coords, coords_are_cartesian=False,
)
frompymatgen.io.cifimportCifWriterwriter=CifWriter(structure, symprec=None)
path.parent.mkdir(parents=True, exist_ok=True)
writer.write_file(str(path))
# ---------------------------------------------------------------------------# 主进程: 调度循环，不加载 GPU 资源# ---------------------------------------------------------------------------defscheduler_main(args):
"""主进程: 逐个启动子进程，每个子进程生成一个晶体后退出。"""output_dir=Path(args.output_dir)
cifs_dir=output_dir/"cifs"cifs_dir.mkdir(parents=True, exist_ok=True)
log_path=output_dir/"generation.log"withopen(log_path, "w") aslog_file:
deflog(msg):
ts=time.strftime("%Y-%m-%d %H:%M:%S")
line=f"[{ts}] {msg}"print(line, flush=True)
log_file.write(line+"\n")
log_file.flush()
log("="*60)
log("PD 晶体生成（独立进程隔离模式）")
log("="*60)
log(f" num_samples = {args.num_samples}")
log(f" dataset = {args.dataset}")
log(f" temperature = {args.temperature}")
log(f" seed = {args.seed}")
log(f" output_dir = {args.output_dir}")
log(f" 每个晶体使用独立子进程，进程退出即释放全部 GPU 显存")
log("="*60)
script_path=Path(__file__).resolve()
t_start=time.time()
successful=0failed=0foridxinrange(args.num_samples):
log(f" [{idx+1}/{args.num_samples}] 启动 worker {idx} ...")
t0=time.time()
try:
proc=subprocess.run(
[
sys.executable, str(script_path),
"--worker", str(idx),
"--output_dir", str(output_dir),
"--dataset", args.dataset,
"--temperature", str(args.temperature),
"--seed", str(args.seed),
"--diffusion_snr", str(args.diffusion_snr),
],
capture_output=True,
text=True,
timeout=600, # 单个晶体最多 10 分钟
)
exceptsubprocess.TimeoutExpired:
failed+=1log(f" 超时 (>600s)")
continueelapsed=time.time() -t0stdout=proc.stdout.strip()
stderr=proc.stderr.strip()
# 解析输出ifproc.returncode==0andf"OK:{idx}"instdout:
successful+=1# 提取原子数信息atom_info=""forlineinstdout.split("\n"):
iff"OK:{idx}"inline:
atom_info=line.split(":", 2)[-1] if":"inlineelse""breaklog(
f" 成功 ({elapsed:.0f}s) {atom_info}"f" [OK={successful} FAIL={failed}]"
)
else:
failed+=1# 提取错误信息fail_reason="unknown"forlineinstdout.split("\n"):
iff"FAIL:{idx}"inline:
fail_reason=line.split(f"FAIL:{idx}:")[-1]
breakifnotfail_reasonorfail_reason=="unknown":
err_tail=stderr[-300:] ifstderrelse""fail_reason=err_tail.replace("\n", " | ")
log(f" 失败 ({elapsed:.0f}s): {fail_reason}")
# 汇总t_elapsed=time.time() -t_startlog("")
log("="*60)
log(f"生成完成!")
log(f" 成功: {successful}/{args.num_samples}")
log(f" 失败: {failed}/{args.num_samples}")
log(f" 总耗时: {t_elapsed:.0f}s ({t_elapsed/60:.1f}min)")
log(f" 平均每晶体: {t_elapsed/max(args.num_samples, 1):.0f}s")
log("="*60)
# 保存元数据cif_files=list(cifs_dir.glob("gen_*.cif"))
meta= {
"framework": "pd",
"dataset": args.dataset,
"num_requested": args.num_samples,
"num_successful": successful,
"num_failed": failed,
"num_saved_cifs": len(cif_files),
"temperature": args.temperature,
"seed": args.seed,
"diffusion_snr": args.diffusion_snr,
"time_seconds": round(t_elapsed, 1),
"avg_per_crystal_seconds": round(
t_elapsed/max(args.num_samples, 1), 1
),
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
}
meta_path=output_dir/"metadata.json"withopen(meta_path, "w") asf:
json.dump(meta, f, indent=2, ensure_ascii=False)
log(f"元数据: {meta_path}")
# ---------------------------------------------------------------------------# 入口# ---------------------------------------------------------------------------defmain():
args=parse_args()
ifargs.workerisnotNone:
# Worker 模式: 在独立子进程中生成单个晶体worker_main(args)
else:
# 调度模式: 逐个启动子进程scheduler_main(args)
if__name__=="__main__":
main()3.3.3 采样指标合并对比移植原版的评估指标Python代码,点击后展开脚本#!/usr/bin/env python3# -*- coding: utf-8 -*-"""SGEquiDiff 生成评估指标模块。从 sgequidiff原版/src/utils/eval_utils.py 移植关键评估函数，"""from __future__ importannotationsimportcopyimportitertoolsimportjsonimportosimportwarningsfromcollectionsimportCounterfrompathlibimportPathfromtypingimportDict, List, Optional, Sequence, TuplefromzipfileimportZipFilefromtempfileimportTemporaryDirectoryimportnumpyasnpfrompymatgen.analysis.structure_matcherimportStructureMatcherfrompymatgen.coreimportComposition, Element, Structure, Latticefrompymatgen.io.aseimportAseAtomsAdaptorfromscipy.spatial.distanceimportjensenshannonfromscipy.statsimportwasserstein_distancefromsmact.screeningimportsmact_filterfromtqdmimporttqdmimportase.io# ============================================================# 常量# ============================================================# MP-20 训练集空间群分布（用于JSD对比）MP20_SPACE_GROUP_DIST=None# MP-20 训练集 Wyckoff 维度分布MP20_WYCKOFF_DIM_DIST=None# MP-20 训练集密度统计MP20_DENSITY_STATS=None# ============================================================# 晶体有效性检查# ============================================================defsmact_validity(composition: Composition) ->bool:
"""SMACT 化学有效性检查。 检查元素的化合价是否能平衡。 """try:
comp_dict=composition.to_reduced_dictelements= [Element(el) forelincomp_dict.keys()]
amounts=list(comp_dict.values())
# smact_filter 返回 (valid, reason)returnsmact_filter(elements, amounts)[0]
exceptException:
returnFalsedefstructure_validity(structure: Structure) ->bool:
"""结构有效性检查。 标准： - 任意两个原子间距 > 0.5 Angstrom - 晶胞体积 > 0.1 Angstrom^3 - 原子数量 > 0 """iflen(structure) ==0:
returnFalseifstructure.volume<0.1:
returnFalse# 检查最小原子间距dist_matrix=structure.distance_matrixnp.fill_diagonal(dist_matrix, np.inf)
min_dist=dist_matrix.min()
ifmin_dist<0.5:
returnFalsereturnTruedefcompute_validity(structures: Sequence[Structure]) ->Dict:
"""计算有效性指标。 Args: structures: pymatgen Structure 列表 Returns: { "total": int, "valid_structure": int, "valid_composition": int, "valid_both": int, "structure_validity_ratio": float, "composition_validity_ratio": float, "validity_ratio": float, } """total=len(structures)
n_struct_valid=0n_comp_valid=0n_valid=0forsinstructures:
struct_ok=structure_validity(s)
comp_ok=smact_validity(s.composition)
ifstruct_ok:
n_struct_valid+=1ifcomp_ok:
n_comp_valid+=1ifstruct_okandcomp_ok:
n_valid+=1return {
"total": total,
"valid_structure": n_struct_valid,
"valid_composition": n_comp_valid,
"valid_both": n_valid,
"structure_validity_ratio": n_struct_valid/max(total, 1),
"composition_validity_ratio": n_comp_valid/max(total, 1),
"validity_ratio": n_valid/max(total, 1),
}
# ============================================================# 结构指纹与多样性# ============================================================defget_structure_fingerprint(structure: Structure) ->np.ndarray:
"""计算结构指纹（加权原子位置直方图）。"""try:
frommatminer.featurizers.structureimportSiteStatsFingerprintfrommatminer.featurizers.siteimportCrystalNNFingerprintnn_fp=CrystalNNFingerprint.from_preset("ops")
site_fp=SiteStatsFingerprint(nn_fp)
returnsite_fp.featurize(structure)
exceptException:
returnnp.zeros(256)
defget_fingerprint_pairwise_dist(fp_array: np.ndarray) ->float:
"""计算指纹矩阵中所有配对之间的平均距离。"""n=len(fp_array)
ifn<=1:
return0.0dists= []
foriinrange(n):
forjinrange(i+1, n):
d=np.linalg.norm(fp_array[i] -fp_array[j])
dists.append(d)
returnfloat(np.mean(dists)) ifdistselse0.0defcompute_diversity(structures: Sequence[Structure]) ->Dict:
"""计算多样性指标。 Returns: { "structural_diversity": float, "composition_diversity": float, } """iflen(structures) <=1:
return {"structural_diversity": 0.0, "composition_diversity": 0.0}
# 结构指纹多样性 - 用固定长度填充处理不一致的指纹维度fingerprints= []
forsintqdm(structures, desc="结构指纹", leave=False):
fp=get_structure_fingerprint(s)
fp=np.asarray(fp).flatten()
fingerprints.append(fp)
# 填充到相同长度max_len=max(len(fp) forfpinfingerprints)
fp_padded=np.zeros((len(fingerprints), max_len))
fori, fpinenumerate(fingerprints):
fp_padded[i, :len(fp)] =fpstruct_div=get_fingerprint_pairwise_dist(fp_padded)
# 组成多样性（基于元素数量向量）comp_vectors= []
forsintqdm(structures, desc="元素向量", leave=False):
comp=s.compositionvec=np.zeros(100)
forel, amtincomp.to_reduced_dict.items():
try:
element=Element(el)
z=element.Zif1<=z<100:
vec[z] =amtexceptException:
continuecomp_vectors.append(vec)
comp_array=np.array(comp_vectors)
comp_div=get_fingerprint_pairwise_dist(comp_array)
return {"structural_diversity": struct_div, "composition_diversity": comp_div}
# ============================================================# 覆盖率（与参考集对比）# ============================================================defcompute_coverage(
gen_structures: Sequence[Structure],
ref_structures: Sequence[Structure],
struc_cutoff: float=0.4,
comp_cutoff: float=10.0,
) ->Dict:
"""计算覆盖率指标（生成结构中有多少比例能在参考集中找到匹配）。 Args: gen_structures: 生成的结构列表 ref_structures: 参考结构列表（如训练集） struc_cutoff: 结构匹配 cutoff comp_cutoff: 组成匹配 cutoff Returns: {"coverage": float, "num_matched": int, "total": int} """matcher=StructureMatcher(
stol=struc_cutoff,
angle_tol=comp_cutoff,
ltol=comp_cutoff,
)
num_matched=0total=len(gen_structures)
# 对每个生成结构，检查是否能在参考集中找到匹配forgen_sintqdm(gen_structures, desc="覆盖率", leave=False):
iftotal==0:
breakmatched=Falseforref_sinref_structures:
ifmatcher.fit(gen_s, ref_s):
matched=Truebreakifmatched:
num_matched+=1return {
"coverage": num_matched/max(total, 1),
"num_matched": num_matched,
"total": total,
}
# ============================================================# 分布指标# ============================================================defget_space_group_number(structure: Structure) ->int:
"""获取空间群编号（1-230）。"""try:
frompymatgen.symmetry.analyzerimportSpacegroupAnalyzersga=SpacegroupAnalyzer(structure)
returnsga.get_space_group_number()
exceptException:
return1defcompute_space_group_jsd(
gen_structures: Sequence[Structure],
ref_distribution: Optional[np.ndarray] =None,
) ->float:
"""计算空间群分布的 Jensen-Shannon Divergence。 Args: gen_structures: 生成结构列表 ref_distribution: 参考分布（230维），默认均匀分布 Returns: JSD 值（越小越接近） """# 计算生成结构的空间群分布sg_counts=Counter()
forsintqdm(gen_structures, desc="空间群", leave=False):
sg=get_space_group_number(s)
sg_counts[sg] +=1gen_dist=np.zeros(230)
forsg, countinsg_counts.items():
if1<=sg<=230:
gen_dist[sg-1] =countgen_dist=gen_dist/max(gen_dist.sum(), 1)
# 参考分布ifref_distributionisnotNone:
ref_dist=np.array(ref_distribution)
else:
ref_dist=np.ones(230) /230# 均匀分布# 计算 JSDreturnfloat(jensenshannon(gen_dist, ref_dist))
defcompute_density_wasserstein(
gen_structures: Sequence[Structure],
ref_densities: Optional[List[float]] =None,
) ->float:
"""计算密度分布的 Wasserstein 距离。"""gen_densities= []
forsintqdm(gen_structures, desc="密度", leave=False):
try:
gen_densities.append(s.density)
exceptException:
passifnotgen_densities:
return1.0ifref_densitiesisnotNoneandlen(ref_densities) >0:
returnfloat(wasserstein_distance(gen_densities, ref_densities))
else:
return0.0defcompute_num_elements_wasserstein(
gen_structures: Sequence[Structure],
ref_num_elements: Optional[List[int]] =None,
) ->float:
"""计算元素数量分布的 Wasserstein 距离。"""gen_counts= []
forsingen_structures:
gen_counts.append(len(s.composition.to_reduced_dict))
ifref_num_elementsisnotNoneandlen(ref_num_elements) >0:
returnfloat(wasserstein_distance(gen_counts, ref_num_elements))
else:
return0.0defcompute_distribution_metrics(
gen_structures: Sequence[Structure],
ref_structures: Optional[Sequence[Structure]] =None,
) ->Dict:
"""计算分布指标。 Returns: { "space_group_jsd": float, "density_wasserstein": float, "num_elements_wasserstein": float, } """# 计算参考分布的统计量ref_sg_dist=Noneref_densities=Noneref_num_elements=Noneifref_structuresisnotNoneandlen(ref_structures) >0:
sg_counts=Counter()
densities= []
num_elems= []
forsintqdm(ref_structures, desc="参考分布", leave=False):
sg=get_space_group_number(s)
sg_counts[sg] +=1try:
densities.append(s.density)
exceptException:
passnum_elems.append(len(s.composition.to_reduced_dict))
ref_sg_dist=np.zeros(230)
forsg, countinsg_counts.items():
if1<=sg<=230:
ref_sg_dist[sg-1] =countref_sg_dist=ref_sg_dist/max(ref_sg_dist.sum(), 1)
ref_densities=densitiesref_num_elements=num_elemsreturn {
"space_group_jsd": compute_space_group_jsd(gen_structures, ref_sg_dist),
"density_wasserstein": compute_density_wasserstein(gen_structures, ref_densities),
"num_elements_wasserstein": compute_num_elements_wasserstein(gen_structures, ref_num_elements),
}
# ============================================================# 结构加载# ============================================================defload_structures_from_npz(path: str) ->List[Structure]:
"""从 .npz 文件加载晶体结构（SGEquiDiff 格式）。"""pass# 需要根据实际 npz 格式实现defload_structures_from_cif_dir(path: str) ->List[Structure]:
"""从 CIF 目录加载结构。"""structures= []
path=Path(path)
ifnotpath.exists():
returnstructuresforfnameinsorted(path.iterdir()):
iffname.suffixin (".cif",):
try:
frompymatgen.io.cifimportCifParserparser=CifParser(str(fname), occupancy_tolerance=100)
s=parser.parse_structures()[0]
structures.append(s)
exceptExceptionase:
warnings.warn(f"读取 {fname.name} 失败: {e}")
returnstructuresdefload_structures_from_json(path: str) ->List[Structure]:
"""从 JSON 文件加载结构。"""structures= []
path=Path(path)
ifnotpath.exists():
returnstructureswithopen(path) asf:
data=json.load(f)
foritemindata:
try:
lattice=Lattice.from_parameters(
*item["lattice_params"]
)
species=item["species"]
coords=item["frac_coords"]
s=Structure(lattice, species, coords, coords_are_cartesian=False)
structures.append(s)
exceptExceptionase:
warnings.warn(f"解析结构失败: {e}")
returnstructures# ============================================================# 完整评估# ============================================================defevaluate_structures(
structures: Sequence[Structure],
ref_structures: Optional[Sequence[Structure]] =None,
dataset_name: str="mp_20",
) ->Dict:
"""对结构列表进行全面的评估。 Args: structures: 要评估的结构列表 ref_structures: 参考结构列表（训练集），用于覆盖率和分布指标 dataset_name: 数据集名称 Returns: 包含所有评估指标的字典 """iflen(structures) ==0:
return {"error": "没有结构可供评估"}
metrics= {}
metrics["num_structures"] =len(structures)
# 1. 有效性print("计算有效性指标...")
metrics["validity"] =compute_validity(structures)
# 2. 多样性print("计算多样性指标...")
metrics["diversity"] =compute_diversity(structures)
# 3. 覆盖率（需要参考集）ifref_structuresisnotNoneandlen(ref_structures) >0:
print("计算覆盖率...")
metrics["coverage"] =compute_coverage(structures, ref_structures)
else:
metrics["coverage"] = {"coverage": 0.0, "num_matched": 0, "total": len(structures)}
# 4. 分布指标print("计算分布指标...")
metrics["distribution"] =compute_distribution_metrics(structures, ref_structures)
returnmetricsdefcompute_comparison(metrics_pt: Dict, metrics_pd: Dict) ->Dict:
"""对比 PT 和 PD 的评估指标，计算相对误差。 Args: metrics_pt: PyTorch 评估结果 metrics_pd: Paddle 评估结果 Returns: { "comparison": {指标名: {"pt": x, "pd": y, "abs_diff": d, "rel_diff": d/|x|}}, "all_within_5pct": bool } """comparison= {}
def_compare(key: str, pt_val: float, pd_val: float):
abs_diff=abs(pt_val-pd_val)
rel_diff=abs_diff/max(abs(pt_val), 1e-10)
return {
"pt": float(pt_val),
"pd": float(pd_val),
"abs_diff": float(abs_diff),
"rel_diff": float(rel_diff),
"within_5pct": rel_diff<0.05,
}
# 有效性指标forkin ("validity_ratio", "structure_validity_ratio", "composition_validity_ratio"):
pt_v=metrics_pt.get("validity", {}).get(k, 0)
pd_v=metrics_pd.get("validity", {}).get(k, 0)
comparison[f"validity.{k}"] =_compare(k, pt_v, pd_v)
# 多样性指标forkin ("structural_diversity", "composition_diversity"):
pt_v=metrics_pt.get("diversity", {}).get(k, 0)
pd_v=metrics_pd.get("diversity", {}).get(k, 0)
comparison[f"diversity.{k}"] =_compare(k, pt_v, pd_v)
# 覆盖率forkin ("coverage",):
pt_v=metrics_pt.get("coverage", {}).get(k, 0)
pd_v=metrics_pd.get("coverage", {}).get(k, 0)
comparison[f"coverage.{k}"] =_compare(k, pt_v, pd_v)
# 分布指标forkin ("space_group_jsd", "density_wasserstein", "num_elements_wasserstein"):
pt_v=metrics_pt.get("distribution", {}).get(k, 0)
pd_v=metrics_pd.get("distribution", {}).get(k, 0)
comparison[f"distribution.{k}"] =_compare(k, pt_v, pd_v)
# 检查是否所有指标都在 5% 以内all_within=all(v.get("within_5pct", False) forvincomparison.values())
return {
"comparison": comparison,
"all_within_5pct": all_within,
"num_metrics": len(comparison),
"num_passed": sum(1forvincomparison.values() ifv.get("within_5pct", False)),
}
defprint_comparison_summary(comparison_result: Dict):
"""打印对比结果摘要。"""print("\n"+"="*70)
print("PT vs PD 采样指标对比报告")
print("="*70)
data=comparison_result["comparison"]
print(f"\n{'指标':<40s}{'PT':>10s}{'PD':>10s}{'|diff|':>12s}{'rel_diff':>10s}{'结果':>8s}")
print("-"*90)
forkey, valsinsorted(data.items()):
status="PASS"ifvals["within_5pct"] else"FAIL"print(
f"{key:<40s}{vals['pt']:>10.4f}{vals['pd']:>10.4f} "f"{vals['abs_diff']:>12.6f}{vals['rel_diff']:>10.4f}{status:>8s}"
)
print("-"*90)
print(f"\n总指标数: {comparison_result['num_metrics']}")
print(f"通过数: {comparison_result['num_passed']}")
print(f"整体结果: {'PASS (全部在 5% 以内)'ifcomparison_result['all_within_5pct'] else'FAIL (部分指标超出 5%)'}")
print("="*70)开始执行对比Python脚本,点击后展开脚本#!/usr/bin/env python3# -*- coding: utf-8 -*-"""晶体采样指标评估与 PT/PD 对比脚本。使用流程: 步骤1 (PT): cd sgequidiff原版代码目录 && source .venv/bin/activate python scripts/generate_samples_pt.py ... 步骤2 (PD): conda activate ppmat python scripts/generate_samples_pd.py ... 步骤3 (对比): 在本脚本中指定 PT/PD 输出目录即可对比"""importargparseimportjsonimportosimportsysimporttimefrompathlibimportPathfromtypingimportDict, List, Optionalimportnumpyasnpdefparse_args():
parser=argparse.ArgumentParser(description="采样指标评估与对比")
# 输入目录parser.add_argument("--pt_cif_dir", type=str, default=None,
help="PT 生成晶体的 CIF 目录")
parser.add_argument("--pd_cif_dir", type=str, default=None,
help="PD 生成晶体的 CIF 目录")
parser.add_argument("--ref_cif_dir", type=str, default=None,
help="参考集 CIF 目录（如训练集）")
# 输出parser.add_argument("--output_dir", type=str, default=None,
help="评估结果输出目录")
parser.add_argument("--dataset", type=str, default="mp_20",
choices=["mp_20", "mpts_52"])
parser.add_argument("--num_ref", type=int, default=None,
help="参考集使用数量，默认全部")
returnparser.parse_args()
defmain():
args=parse_args()
# 确定输出目录ifargs.output_dirisNone:
timestamp=time.strftime("%Y%m%d_%H%M%S")
output_dir=Path(f"./outputs/sampling_eval_{timestamp}")
else:
output_dir=Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
# 日志log_path=output_dir/"evaluation.log"log_file=open(log_path, "w")
deflog(msg):
print(msg, flush=True)
log_file.write(msg+"\n")
log_file.flush()
log(f"{'='*70}")
log(f"SGEquiDiff 采样指标评估")
log(f"{'='*70}")
log(f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
log(f"数据集: {args.dataset}")
log(f"输出目录: {output_dir}")
log("")
# 导入评估模块sys.path.insert(0, str(Path(__file__).resolve().parent))
fromgeneration_metricsimport (
load_structures_from_cif_dir,
load_structures_from_json,
evaluate_structures,
compute_comparison,
print_comparison_summary,
)
# 加载 PT 结构pt_structures= []
ifargs.pt_cif_dir:
pt_path=Path(args.pt_cif_dir)
ifpt_path.exists():
log(f"加载 PT 结构: {pt_path}")
pt_structures=load_structures_from_cif_dir(str(pt_path))
log(f" PT 结构数: {len(pt_structures)}")
else:
log(f" PT 目录不存在: {pt_path}")
# 加载 PD 结构pd_structures= []
ifargs.pd_cif_dir:
pd_path=Path(args.pd_cif_dir)
ifpd_path.exists():
log(f"加载 PD 结构: {pd_path}")
pd_structures=load_structures_from_cif_dir(str(pd_path))
log(f" PD 结构数: {len(pd_structures)}")
else:
log(f" PD 目录不存在: {pd_path}")
# 加载参考结构ref_structures= []
ifargs.ref_cif_dir:
ref_path=Path(args.ref_cif_dir)
ifref_path.exists():
log(f"加载参考结构: {ref_path}")
all_ref=load_structures_from_cif_dir(str(ref_path))
ifargs.num_ref:
ref_structures=all_ref[:args.num_ref]
else:
ref_structures=all_reflog(f" 参考结构数: {len(ref_structures)}")
else:
log(f" 参考目录不存在: {ref_path}")
# 评估 PTmetrics_pt=Noneifpt_structures:
log(f"\n{'='*70}")
log("评估 PT 生成样本...")
log(f"{'='*70}")
metrics_pt=evaluate_structures(pt_structures, ref_structures, args.dataset)
_save_metrics(output_dir/"metrics_pt.json", metrics_pt)
log(f"\nPT 评估结果已保存到: {output_dir/'metrics_pt.json'}")
else:
log("\n跳过 PT 评估 (无结构)")
# 评估 PDmetrics_pd=Noneifpd_structures:
log(f"\n{'='*70}")
log("评估 PD 生成样本...")
log(f"{'='*70}")
metrics_pd=evaluate_structures(pd_structures, ref_structures, args.dataset)
_save_metrics(output_dir/"metrics_pd.json", metrics_pd)
log(f"\nPD 评估结果已保存到: {output_dir/'metrics_pd.json'}")
else:
log("\n跳过 PD 评估 (无结构)")
# 对比comparison=Noneifmetrics_ptandmetrics_pd:
log(f"\n{'='*70}")
log("PT vs PD 对比")
log(f"{'='*70}")
comparison=compute_comparison(metrics_pt, metrics_pd)
print_comparison_summary(comparison)
# 保存对比结果withopen(output_dir/"comparison.json", "w") asf:
# 转换 numpy 类型json.dump(comparison, f, indent=2, default=_json_default)
log(f"\n对比结果已保存到: {output_dir/'comparison.json'}")
elifpt_structures:
log("\n仅 PT 评估完成 (无 PD 数据无法对比)")
elifpd_structures:
log("\n仅 PD 评估完成 (无 PT 数据无法对比)")
else:
log("\n未加载任何结构，评估跳过")
# 汇总log(f"\n{'='*70}")
log("评估汇总")
log(f"{'='*70}")
ifcomparison:
log(f" 总指标数: {comparison['num_metrics']}")
log(f" 通过数: {comparison['num_passed']}")
log(f" 结果: {'全部通过'ifcomparison['all_within_5pct'] else'部分未通过'}")
log(f" 输出目录: {output_dir}")
log(f" 日志文件: {log_path}")
log(f"{'='*70}")
log_file.close()
def_save_metrics(path, metrics):
"""保存评估指标到 JSON。"""withopen(path, "w") asf:
json.dump(metrics, f, indent=2, default=_json_default)
def_json_default(obj):
"""JSON 序列化的默认处理器。"""ifisinstance(obj, (np.integer,)):
returnint(obj)
ifisinstance(obj, (np.floating,)):
returnfloat(obj)
ifisinstance(obj, (np.ndarray,)):
returnobj.tolist()
returnstr(obj)
if__name__=="__main__":
main() |
leeleolay
requested changes
Jun 24, 2026
| --- | ||
| ## Command |
| | Dataset | Sub-module | Download | | ||
| | --- | --- | --- | | ||
| | mp_20 | diffusion | [download](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mp_20_best_diffusion_snapshot.pdparams) | |
| --- | ||
| ## Results |
Collaborator
There was a problem hiding this comment.
不符合已有dataset格式，另外为什么要新建asu单独的dataset
Collaborator
There was a problem hiding this comment.
resources下的文件先建议都放到ppmat/utils/vocabs/crystals/下面
Comment on lines
+66
to
+79
| PRETRAINED_WEIGHT_URLS = { | ||
| "mp_20": { | ||
| "diffusion": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mp_20_best_diffusion_snapshot.pdparams", | ||
| "lattice": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mp_20_best_lattice_snapshot.pdparams", | ||
| "space_group": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mp_20_best_space_group_snapshot.pdparams", | ||
| "wyckoff": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mp_20_best_wyckoff-transformer_snapshot.pdparams", | ||
| }, | ||
| "mpts_52": { | ||
| "diffusion": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mpts_52_best_diffusion_snapshot.pdparams", | ||
| "lattice": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mpts_52_best_lattice_snapshot.pdparams", | ||
| "space_group": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mpts_52_best_space_group_snapshot.pdparams", | ||
| "wyckoff": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/SGEquiDiff/mpts_52_best_wyckoff-transformer_snapshot.pdparams", | ||
| }, | ||
| } |
Collaborator
There was a problem hiding this comment.
放到ppmat/models/init 里的registry，且支持一键推理/采样
Collaborator
There was a problem hiding this comment.
基于已有的sampler和scheduler功能，需重构相关diffusion实现
learncat163force-pushed
the
feature/add-sgequidiff
branch
from
June 29, 2026 08:58
1f37fb0 to
f1ed062Comparelearncat163
commented
Jul 2, 2026
Author
放弃了之前的 1:1的代码转译的paconvert的代码模式； 重新组织和复用了代码。 |
Thanks for your contribution! |
leeleolay
commented
Jul 16, 2026
Collaborator
存在冲突 |
leeleolay
left a comment
Collaborator
There was a problem hiding this comment.
还是不符合格式要求，model的部分和dataset的部分辛苦重点整理，model需满足单模型文件策略，与model的无关的部分均在data里实现并使用复用相关已有基础设施
Collaborator
There was a problem hiding this comment.
没有按照已有的规范，参考mp20格式，可复制去代码，来修改，注意build的逻辑，cache的逻辑
…uidiff # Conflicts: # ppmat/datasets/__init__.py # ppmat/models/__init__.py
…uidiff # Conflicts: # ppmat/utils/crystal.py
learncat163
commented
Aug 19, 2026
Author
@leeleolay 重构了一版 |
…uidiff Resolve ppmat/utils/scatter.py by keeping functions from both sides: scatter_argmax/scatter_min_with_argmin/scatter_min_indices (HEAD) and scatter_sum_first_order/scatter_mean/scatter_min (upstream). Fix scatter_sum body to use the upstream corrected call. Update sgequidiff sample.py docs from --output_dir to --output_path. Co-authored-by: CommandCodeBot <noreply@commandcode.ai>
leeleolay
left a comment
Collaborator
There was a problem hiding this comment.
vocab目前套件内实现了一个自动加载和配置的文件，辛苦参考下，相关需要我上传的vocab 辛苦发给我
Comment on lines
+165
to
+178
| ## Environment Requirements | ||
| The model is developed and tested under the following environment: | ||
| | Package | Version | | ||
| | --- | --- | | ||
| | Python | 3.10 | | ||
| | PaddlePaddle | >= 3.1 | | ||
| | paddle_scatter | from source | | ||
| | pymatgen | 2024.10.29 | | ||
| | scipy | 1.13.1 | | ||
| | numpy | 1.26.4 | | ||
| Refer to [Install.md](../../../Install.md) for the complete installation instructions. |
Comment on lines
+184
to
+190
| ### Setup | ||
| NPZ data is auto-discovered in this order: `$ASU_DATA_DIR` > `data/data/` > `~/.asu_data/`. Place `train.npz / val.npz / test.npz` under `<root>/mp_20/`. Set the env var only if data lives elsewhere: | ||
| ```bash | ||
| export ASU_DATA_DIR=/path/to/data | ||
| ``` |
Comment on lines
+199
to
+214
| ### Key Configuration | ||
| Key hyperparameters shared by the released configs (`sgequidiff_mp20.yaml` / `sgequidiff_mpts_52.yaml`): | ||
| | Parameter | Value | Description | | ||
| | --- | --- | --- | | ||
| | `num_timesteps` | 1000 | diffusion timesteps | | ||
| | `sigma_min` / `sigma_max` | 0.002 / 0.5 | VE-SDE noise schedule bounds | | ||
| | `noise_scheduler_num_monte_carlo_samples` | 2500 | MC samples for the ASU-wrapped sigma-norm table | | ||
| | `num_wn_lattice_translations` | 3 | wrapped-normal lattice translations | | ||
| | `model_type` | `gnn` | non-equivariant backbone (`mlp` / `gnn` / `cspnet`) | | ||
| | `time_emb_dim` | 128 | Fourier time embedding dimension | | ||
| | `batch_size` | 32 (train) / 64 (val, test) | dataloader batch size | | ||
| | optimizer | AdamW, lr 1e-3 | ReduceOnPlateau (factor 0.6, patience 30, min_lr 1e-5) | | ||
| | `max_epochs` | 2000 | total training epochs | | ||
Comment on lines
+192
to
+197
| ### Runtime caches | ||
| On the first run the model generates two cache files under the same data root as the NPZ files (`$ASU_DATA_DIR` > `data/data/` > `~/.asu_data/`): | ||
| - `wyckoff_shape_decomposition.pkl`: precomputed Wyckoff site shapes, built from the bundled JSON in `ppmat/models/sgequidiff/vocabs/`. | ||
| - `expected_score_norms_*.pdparams`: Monte Carlo score-norm tables over space groups / Wyckoff sites / timesteps; the filename encodes `sigma_min/max`, `num_timesteps`, MC samples and lattice translations. First computation is heavy (230 space groups); later runs with the same diffusion schedule reuse the file. |
leeleolay
left a comment
Collaborator
There was a problem hiding this comment.
辛苦参考相关意见重新整理下代码架构，model内的部分辛苦再整理精简下
| return symm_lattice_matrix | ||
| def build_pbc_graph( |
Collaborator
There was a problem hiding this comment.
这个看起来是构建graph，相关功能应通过build graph来实现，保存位置应在ppmat/models/common/graph_converter.py
Comment on lines
+386
to
+399
| named_lr_groups (Optional[List[Dict[str, Union[str, float]]]], optional): | ||
| Per-layer learning-rate (and optionally weight-decay) multipliers. | ||
| Each entry is a dict like ``{"name": "layer_segment", | ||
| "lr_multiplier": 0.1, "weight_decay_multiplier": 0.0}``. The | ||
| ``name`` field matches any dot-separated segment of a parameter's | ||
| full path (avoiding substring false positives such as | ||
| ``"time_embedder"`` accidentally matching | ||
| ``"time_embedder_aux.weight"``); the first matching entry wins. | ||
| The effective learning rate of a matched parameter is | ||
| ``learning_rate * lr_multiplier``; if ``weight_decay_multiplier`` | ||
| is given, the effective weight decay is also scaled | ||
| (``weight_decay * weight_decay_multiplier``). Unmatched parameters | ||
| fall back to ``lr_multiplier=1.0`` and no weight-decay multiplier. | ||
| Defaults to None. |
Collaborator
There was a problem hiding this comment.
vocab的功能辛苦参考下套件内的实现和load的方式，如果有关于vocab实现的规范性问题，可以讨论
Comment on lines
+45
to
+62
| @dataclasses.dataclass | ||
| class EquivariantDiffusionModelConfig: | ||
| """Diffusion model hyperparameter config.""" | ||
| num_lattice_translations: int = 3 | ||
| noise_scheduler_num_monte_carlo_samples: int = 2_500 | ||
| num_timesteps: int = 1000 | ||
| sigma_min: float = 0.002 | ||
| sigma_max: float = 0.5 | ||
| time_emb_dim: int = 128 | ||
| model_type: str = "gnn" # ["mlp", "gnn", "cspnet"] | ||
| num_plane_wave_freqs: int = 96 | ||
| mlp_hidden_dim: int = 128 | ||
| gnn_config: Optional[GNNConfig] = None | ||
| cspnet_config: Optional[CSPNetConfig] = None | ||
| noise_scheduler_cfg: Optional[dict] = None | ||
leeleolay
left a comment
Collaborator
There was a problem hiding this comment.
Dataset 改为显式 path，删除项目根目录和 ~/.asu_data 推断
接入统一下载、MD5 和缓存机制
修复 unconditional 模型与 by_num_atoms/formula API 的冲突
Comment on lines
+33
to
+36
| _CANDIDATE_DIRS = [ | ||
| _PROJECT_ROOT / "data/data", | ||
| Path("~").expanduser() / ".asu_data", | ||
| ] |
| from ppmat.datasets.custom_data_type import ConcatData | ||
| from ppmat.models.sgequidiff.sgequidiff_meta import ELEMENT_ENCODING_SIZE as _NUM_ELEMENTS | ||
| _PROJECT_ROOT = Path(__file__).resolve().parents[2] |
Comment on lines
+39
to
+64
| def resolve_asu_data_dir() -> Path: | ||
| """Locate the ASU data directory. | ||
| Priority: $ASU_DATA_DIR > project data dirs > ~/.asu_data. | ||
| A directory is accepted if it contains any of the supported ASU | ||
| dataset directories (npz files). Model-specific resources (e.g. | ||
| ``ppmat.models.sgequidiff.vocabs``) live in their own module and | ||
| resolve via their own resolvers rather than this directory. | ||
| """ | ||
| env = os.getenv("ASU_DATA_DIR") | ||
| if env: | ||
| return Path(env) | ||
| for candidate in _CANDIDATE_DIRS: | ||
| if any((candidate / d).is_dir() for d in SUPPORTED_DATASETS): | ||
| return candidate | ||
| _candidate_paths = "\n".join([f" * {d}" for d in _CANDIDATE_DIRS]) | ||
| raise FileNotFoundError( | ||
| f"Cannot find ASU data directory.\n" | ||
| f"\nTried paths:\n{_candidate_paths}\n" | ||
| f"\nSolutions:\n" | ||
| f" 1. Copy data to: {str(_CANDIDATE_DIRS[0])}\n" | ||
| f" 2. Set environment variable: export ASU_DATA_DIR=/your/data/path\n" | ||
| f"\nRequired: at least one of dataset dirs {list(SUPPORTED_DATASETS)}" | ||
| ) |
| ) | ||
| # Supported data splits for ASU datasets (each has a corresponding NPZ archive). | ||
| SUPPORTED_SPLITS = ("train", "val", "test") |
Comment on lines
+87
to
+102
| def _parse_flat(flat: np.ndarray): | ||
| """Parse one flat NPZ crystal array into its fields (see layout above).""" | ||
| n = int(flat[0]) | ||
| sg = int(flat[_IDX_SG]) - 1 | ||
| comp = flat[_IDX_COMP:_IDX_LENGTHS].astype(np.float32) | ||
| lengths = flat[_IDX_LENGTHS:_IDX_ANGLES].astype(np.float32) | ||
| angles = flat[_IDX_ANGLES:_IDX_ATOMS].astype(np.float32) | ||
| elems = flat[_IDX_ATOMS : _IDX_ATOMS + n].astype(np.int64) | ||
| wycks = flat[_IDX_ATOMS + n : _IDX_ATOMS + 2 * n].astype(np.int64) | ||
| fracs = ( | ||
| flat[_IDX_ATOMS + 2 * n : _IDX_ATOMS + 5 * n].reshape(n, 3).astype(np.float32) | ||
| ) | ||
| wsi = None | ||
| if len(flat) > _IDX_ATOMS + 5 * n: | ||
| wsi = flat[_IDX_ATOMS + 5 * n : _IDX_ATOMS + 6 * n].astype(np.int64) | ||
| return sg, comp, lengths, angles, n, elems, wycks, fracs, wsi |
Comment on lines
+112
to
+126
| def __init__( | ||
| self, | ||
| name: str = "mp_20", | ||
| split: str = "train", | ||
| data_directory: Optional[Path] = None, | ||
| ): | ||
| super().__init__() | ||
| assert split in SUPPORTED_SPLITS, f"unknown split: {split}" | ||
| assert name in SUPPORTED_DATASETS, f"unknown name: {name}" | ||
| if data_directory is None: | ||
| data_directory = resolve_asu_data_dir() | ||
| npz = np.load(Path(data_directory) / name / f"{split}.npz") | ||
| self._flat_crystals = np.split(npz["packed"], npz["indices"]) |
Collaborator
There was a problem hiding this comment.
看下mp20初始化函数和各方法，保证基本的方法名一致，行为保持一致
| } | ||
| def build_pbc_graph( |
Collaborator
There was a problem hiding this comment.
build_graph 方法内有实现pgl封装，当前实现方式不符合已有方式
| class ASUCrystal: | ||
| """ASU crystal data.""" | ||
| __hash__ = None |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
来自废弃的PR #287
已经解决的有评审意见