【MIIT program】Feature add SGEquiDiff - #300
Open
learncat163 wants to merge 2 commits into
Open
Conversation
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
for t_val in 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)
for snapshot_path in 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 = {}
for t_val in 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 = []
for key in sorted(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 for r in all_rows if r[0].startswith("diffusion.")]
other_rows = [r for r in all_rows if not r[0].startswith("diffusion.")]
if diffusion_rows:
print("\n--- Diffusion Model (predict_equivariant_vectors) ---")
for path, max_diff, mean_diff, status, is_stoch in diffusion_rows:
t_val = path.split(".")[-1]
print(f" t={t_val}: max|diff|={max_diff:.2e} mean|diff|={mean_diff:.2e} [{status}]")
for module_name in ["space_group", "lattice", "wyckoff"]:
module_rows = [r for r in other_rows if r[0].startswith(f"{module_name}.")]
if module_rows:
print(f"\n--- {module_name.replace('_', ' ').title()} ---")
for path, max_diff, mean_diff, status, is_stoch in 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 for r in all_rows if not r[4] and r[1] is not None]
stochastic_rows = [r for r in all_rows if r[4]]
if deterministic_rows:
det_max_diffs = [r[1] for r in 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) for n, p in 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)
for snapshot_path in 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 _, p in 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
for epoch in range(1, args.epochs + 1):
for step_in_epoch, t_val in 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
for p in dm.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 for n, _ 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"),
]
for name, filename in weight_files:
weight_path = ckpt_dir / filename
if weight_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 model
if 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 _, p in 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
for epoch in range(1, args.epochs + 1):
for step_in_epoch, t_val in 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
for p in dm.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 for n, _ 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 = []
for a, b in 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段！！！ |
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
"""
import argparse
import json
import os
import sys
import time
from pathlib import Path
def parse_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"])
return parser.parse_args()
def main():
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.stdout
print(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))
# 导入原始模块
from scripts.generate_crystals import main as raw_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"),
]
if args.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"),
}
with open(output_dir / "metadata.json", "w") as f:
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
"""
import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path
def parse_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 个晶体后退出")
return parser.parse_args()
# ---------------------------------------------------------------------------
# Worker: 在独立子进程中运行，生成单个晶体后退出，GPU 显存随之彻底释放
# ---------------------------------------------------------------------------
def worker_main(args):
"""子进程入口: 构建模型 -> 采样 -> 保存 CIF -> 退出。"""
project_root = str(Path(__file__).resolve().parents[4])
if project_root not in sys.path:
sys.path.insert(0, project_root)
import paddle
from ppmat.models.sgequidiff.weight_utils import load_pretrained_weights
from ppmat.models.sgequidiff.crystal_sampler import (
CrystalSampler, CrystalSamplerConfig,
)
from ppmat.models.sgequidiff.diffusion_model import (
EquivariantDiffusionModelConfig,
)
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.constants import (
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,
)
except Exception as e:
print(f"FAIL:{idx}:sample_error:{e}", flush=True)
return
if not crystals or len(crystals) == 0:
print(f"FAIL:{idx}:empty_result", flush=True)
return
# 保存 CIF
c = crystals[0]
lattice = (
c.conventional_lattice_lengths.tolist()
+ c.conventional_lattice_angles.tolist()
)
elements = [chemical_symbols[i] for i in c.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(1 for e in elements if e not in ("X", "X0+", ""))
print(f"OK:{idx}:atoms={real_atoms}", flush=True)
# 显式清理（虽然进程即将退出，但保险起见）
del model, crystals
paddle.device.cuda.empty_cache()
# ---------------------------------------------------------------------------
# CIF 保存（混合占位模式，与 PT 版本对齐）
# ---------------------------------------------------------------------------
def _save_cif(path, lattice_params, elements, frac_coords, space_group):
"""将晶体信息保存为 CIF 文件。
相同分数坐标的原子合并为混合占位 (mixed-occupancy) 位点，
使 CIF 的位点数与 PT 版本一致，消除评估偏差。
"""
from collections import Counter, defaultdict
from pymatgen.core import Element, Lattice, Structure
# 过滤占位元素
valid_pairs = []
for i, el_name in enumerate(elements):
if el_name in ("X", "X0+", "") or len(el_name) == 0:
continue
try:
el = Element(el_name)
coord = tuple(frac_coords[i])
valid_pairs.append((el, coord))
except Exception:
continue
if not valid_pairs:
return
# 按坐标分组，合并相同坐标的不同元素为混合占位
site_map = defaultdict(list)
for el, coord in valid_pairs:
site_map[coord].append(el)
group_species = []
group_coords = []
for coord, elements_list in site_map.items():
el_counts = Counter(elements_list)
total = sum(el_counts.values())
if total == 1:
group_species.append(elements_list[0])
else:
group_species.append(
{el: count / total for el, count in el_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,
)
from pymatgen.io.cif import CifWriter
writer = CifWriter(structure, symprec=None)
path.parent.mkdir(parents=True, exist_ok=True)
writer.write_file(str(path))
# ---------------------------------------------------------------------------
# 主进程: 调度循环，不加载 GPU 资源
# ---------------------------------------------------------------------------
def scheduler_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"
with open(log_path, "w") as log_file:
def log(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 = 0
failed = 0
for idx in range(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 分钟
)
except subprocess.TimeoutExpired:
failed += 1
log(f" 超时 (>600s)")
continue
elapsed = time.time() - t0
stdout = proc.stdout.strip()
stderr = proc.stderr.strip()
# 解析输出
if proc.returncode == 0 and f"OK:{idx}" in stdout:
successful += 1
# 提取原子数信息
atom_info = ""
for line in stdout.split("\n"):
if f"OK:{idx}" in line:
atom_info = line.split(":", 2)[-1] if ":" in line else ""
break
log(
f" 成功 ({elapsed:.0f}s) {atom_info}"
f" [OK={successful} FAIL={failed}]"
)
else:
failed += 1
# 提取错误信息
fail_reason = "unknown"
for line in stdout.split("\n"):
if f"FAIL:{idx}" in line:
fail_reason = line.split(f"FAIL:{idx}:")[-1]
break
if not fail_reason or fail_reason == "unknown":
err_tail = stderr[-300:] if stderr else ""
fail_reason = err_tail.replace("\n", " | ")
log(f" 失败 ({elapsed:.0f}s): {fail_reason}")
# 汇总
t_elapsed = time.time() - t_start
log("")
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"
with open(meta_path, "w") as f:
json.dump(meta, f, indent=2, ensure_ascii=False)
log(f"元数据: {meta_path}")
# ---------------------------------------------------------------------------
# 入口
# ---------------------------------------------------------------------------
def main():
args = parse_args()
if args.worker is not None:
# 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__ import annotations
import copy
import itertools
import json
import os
import warnings
from collections import Counter
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple
from zipfile import ZipFile
from tempfile import TemporaryDirectory
import numpy as np
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.core import Composition, Element, Structure, Lattice
from pymatgen.io.ase import AseAtomsAdaptor
from scipy.spatial.distance import jensenshannon
from scipy.stats import wasserstein_distance
from smact.screening import smact_filter
from tqdm import tqdm
import ase.io
# ============================================================
# 常量
# ============================================================
# MP-20 训练集空间群分布（用于JSD对比）
MP20_SPACE_GROUP_DIST = None
# MP-20 训练集 Wyckoff 维度分布
MP20_WYCKOFF_DIM_DIST = None
# MP-20 训练集密度统计
MP20_DENSITY_STATS = None
# ============================================================
# 晶体有效性检查
# ============================================================
def smact_validity(composition: Composition) -> bool:
"""SMACT 化学有效性检查。
检查元素的化合价是否能平衡。
"""
try:
comp_dict = composition.to_reduced_dict
elements = [Element(el) for el in comp_dict.keys()]
amounts = list(comp_dict.values())
# smact_filter 返回 (valid, reason)
return smact_filter(elements, amounts)[0]
except Exception:
return False
def structure_validity(structure: Structure) -> bool:
"""结构有效性检查。
标准：
- 任意两个原子间距 > 0.5 Angstrom
- 晶胞体积 > 0.1 Angstrom^3
- 原子数量 > 0
"""
if len(structure) == 0:
return False
if structure.volume < 0.1:
return False
# 检查最小原子间距
dist_matrix = structure.distance_matrix
np.fill_diagonal(dist_matrix, np.inf)
min_dist = dist_matrix.min()
if min_dist < 0.5:
return False
return True
def compute_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 = 0
n_comp_valid = 0
n_valid = 0
for s in structures:
struct_ok = structure_validity(s)
comp_ok = smact_validity(s.composition)
if struct_ok:
n_struct_valid += 1
if comp_ok:
n_comp_valid += 1
if struct_ok and comp_ok:
n_valid += 1
return {
"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),
}
# ============================================================
# 结构指纹与多样性
# ============================================================
def get_structure_fingerprint(structure: Structure) -> np.ndarray:
"""计算结构指纹（加权原子位置直方图）。"""
try:
from matminer.featurizers.structure import SiteStatsFingerprint
from matminer.featurizers.site import CrystalNNFingerprint
nn_fp = CrystalNNFingerprint.from_preset("ops")
site_fp = SiteStatsFingerprint(nn_fp)
return site_fp.featurize(structure)
except Exception:
return np.zeros(256)
def get_fingerprint_pairwise_dist(fp_array: np.ndarray) -> float:
"""计算指纹矩阵中所有配对之间的平均距离。"""
n = len(fp_array)
if n <= 1:
return 0.0
dists = []
for i in range(n):
for j in range(i + 1, n):
d = np.linalg.norm(fp_array[i] - fp_array[j])
dists.append(d)
return float(np.mean(dists)) if dists else 0.0
def compute_diversity(structures: Sequence[Structure]) -> Dict:
"""计算多样性指标。
Returns:
{
"structural_diversity": float,
"composition_diversity": float,
}
"""
if len(structures) <= 1:
return {"structural_diversity": 0.0, "composition_diversity": 0.0}
# 结构指纹多样性 - 用固定长度填充处理不一致的指纹维度
fingerprints = []
for s in tqdm(structures, desc="结构指纹", leave=False):
fp = get_structure_fingerprint(s)
fp = np.asarray(fp).flatten()
fingerprints.append(fp)
# 填充到相同长度
max_len = max(len(fp) for fp in fingerprints)
fp_padded = np.zeros((len(fingerprints), max_len))
for i, fp in enumerate(fingerprints):
fp_padded[i, :len(fp)] = fp
struct_div = get_fingerprint_pairwise_dist(fp_padded)
# 组成多样性（基于元素数量向量）
comp_vectors = []
for s in tqdm(structures, desc="元素向量", leave=False):
comp = s.composition
vec = np.zeros(100)
for el, amt in comp.to_reduced_dict.items():
try:
element = Element(el)
z = element.Z
if 1 <= z < 100:
vec[z] = amt
except Exception:
continue
comp_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}
# ============================================================
# 覆盖率（与参考集对比）
# ============================================================
def compute_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 = 0
total = len(gen_structures)
# 对每个生成结构，检查是否能在参考集中找到匹配
for gen_s in tqdm(gen_structures, desc="覆盖率", leave=False):
if total == 0:
break
matched = False
for ref_s in ref_structures:
if matcher.fit(gen_s, ref_s):
matched = True
break
if matched:
num_matched += 1
return {
"coverage": num_matched / max(total, 1),
"num_matched": num_matched,
"total": total,
}
# ============================================================
# 分布指标
# ============================================================
def get_space_group_number(structure: Structure) -> int:
"""获取空间群编号（1-230）。"""
try:
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
sga = SpacegroupAnalyzer(structure)
return sga.get_space_group_number()
except Exception:
return 1
def compute_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()
for s in tqdm(gen_structures, desc="空间群", leave=False):
sg = get_space_group_number(s)
sg_counts[sg] += 1
gen_dist = np.zeros(230)
for sg, count in sg_counts.items():
if 1 <= sg <= 230:
gen_dist[sg - 1] = count
gen_dist = gen_dist / max(gen_dist.sum(), 1)
# 参考分布
if ref_distribution is not None:
ref_dist = np.array(ref_distribution)
else:
ref_dist = np.ones(230) / 230 # 均匀分布
# 计算 JSD
return float(jensenshannon(gen_dist, ref_dist))
def compute_density_wasserstein(
gen_structures: Sequence[Structure],
ref_densities: Optional[List[float]] = None,
) -> float:
"""计算密度分布的 Wasserstein 距离。"""
gen_densities = []
for s in tqdm(gen_structures, desc="密度", leave=False):
try:
gen_densities.append(s.density)
except Exception:
pass
if not gen_densities:
return 1.0
if ref_densities is not None and len(ref_densities) > 0:
return float(wasserstein_distance(gen_densities, ref_densities))
else:
return 0.0
def compute_num_elements_wasserstein(
gen_structures: Sequence[Structure],
ref_num_elements: Optional[List[int]] = None,
) -> float:
"""计算元素数量分布的 Wasserstein 距离。"""
gen_counts = []
for s in gen_structures:
gen_counts.append(len(s.composition.to_reduced_dict))
if ref_num_elements is not None and len(ref_num_elements) > 0:
return float(wasserstein_distance(gen_counts, ref_num_elements))
else:
return 0.0
def compute_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 = None
ref_densities = None
ref_num_elements = None
if ref_structures is not None and len(ref_structures) > 0:
sg_counts = Counter()
densities = []
num_elems = []
for s in tqdm(ref_structures, desc="参考分布", leave=False):
sg = get_space_group_number(s)
sg_counts[sg] += 1
try:
densities.append(s.density)
except Exception:
pass
num_elems.append(len(s.composition.to_reduced_dict))
ref_sg_dist = np.zeros(230)
for sg, count in sg_counts.items():
if 1 <= sg <= 230:
ref_sg_dist[sg - 1] = count
ref_sg_dist = ref_sg_dist / max(ref_sg_dist.sum(), 1)
ref_densities = densities
ref_num_elements = num_elems
return {
"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),
}
# ============================================================
# 结构加载
# ============================================================
def load_structures_from_npz(path: str) -> List[Structure]:
"""从 .npz 文件加载晶体结构（SGEquiDiff 格式）。"""
pass # 需要根据实际 npz 格式实现
def load_structures_from_cif_dir(path: str) -> List[Structure]:
"""从 CIF 目录加载结构。"""
structures = []
path = Path(path)
if not path.exists():
return structures
for fname in sorted(path.iterdir()):
if fname.suffix in (".cif",):
try:
from pymatgen.io.cif import CifParser
parser = CifParser(str(fname), occupancy_tolerance=100)
s = parser.parse_structures()[0]
structures.append(s)
except Exception as e:
warnings.warn(f"读取 {fname.name} 失败: {e}")
return structures
def load_structures_from_json(path: str) -> List[Structure]:
"""从 JSON 文件加载结构。"""
structures = []
path = Path(path)
if not path.exists():
return structures
with open(path) as f:
data = json.load(f)
for item in data:
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)
except Exception as e:
warnings.warn(f"解析结构失败: {e}")
return structures
# ============================================================
# 完整评估
# ============================================================
def evaluate_structures(
structures: Sequence[Structure],
ref_structures: Optional[Sequence[Structure]] = None,
dataset_name: str = "mp_20",
) -> Dict:
"""对结构列表进行全面的评估。
Args:
structures: 要评估的结构列表
ref_structures: 参考结构列表（训练集），用于覆盖率和分布指标
dataset_name: 数据集名称
Returns:
包含所有评估指标的字典
"""
if len(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. 覆盖率（需要参考集）
if ref_structures is not None and len(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)
return metrics
def compute_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,
}
# 有效性指标
for k in ("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)
# 多样性指标
for k in ("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)
# 覆盖率
for k in ("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)
# 分布指标
for k in ("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) for v in comparison.values())
return {
"comparison": comparison,
"all_within_5pct": all_within,
"num_metrics": len(comparison),
"num_passed": sum(1 for v in comparison.values() if v.get("within_5pct", False)),
}
def print_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)
for key, vals in sorted(data.items()):
status = "PASS" if vals["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% 以内)' if comparison_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 输出目录即可对比
"""
import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional
import numpy as np
def parse_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="参考集使用数量，默认全部")
return parser.parse_args()
def main():
args = parse_args()
# 确定输出目录
if args.output_dir is None:
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")
def log(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))
from generation_metrics import (
load_structures_from_cif_dir,
load_structures_from_json,
evaluate_structures,
compute_comparison,
print_comparison_summary,
)
# 加载 PT 结构
pt_structures = []
if args.pt_cif_dir:
pt_path = Path(args.pt_cif_dir)
if pt_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 = []
if args.pd_cif_dir:
pd_path = Path(args.pd_cif_dir)
if pd_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 = []
if args.ref_cif_dir:
ref_path = Path(args.ref_cif_dir)
if ref_path.exists():
log(f"加载参考结构: {ref_path}")
all_ref = load_structures_from_cif_dir(str(ref_path))
if args.num_ref:
ref_structures = all_ref[:args.num_ref]
else:
ref_structures = all_ref
log(f" 参考结构数: {len(ref_structures)}")
else:
log(f" 参考目录不存在: {ref_path}")
# 评估 PT
metrics_pt = None
if pt_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 评估 (无结构)")
# 评估 PD
metrics_pd = None
if pd_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 = None
if metrics_pt and metrics_pd:
log(f"\n{'='*70}")
log("PT vs PD 对比")
log(f"{'='*70}")
comparison = compute_comparison(metrics_pt, metrics_pd)
print_comparison_summary(comparison)
# 保存对比结果
with open(output_dir / "comparison.json", "w") as f:
# 转换 numpy 类型
json.dump(comparison, f, indent=2, default=_json_default)
log(f"\n对比结果已保存到: {output_dir / 'comparison.json'}")
elif pt_structures:
log("\n仅 PT 评估完成 (无 PD 数据无法对比)")
elif pd_structures:
log("\n仅 PD 评估完成 (无 PT 数据无法对比)")
else:
log("\n未加载任何结构，评估跳过")
# 汇总
log(f"\n{'='*70}")
log("评估汇总")
log(f"{'='*70}")
if comparison:
log(f" 总指标数: {comparison['num_metrics']}")
log(f" 通过数: {comparison['num_passed']}")
log(f" 结果: {'全部通过' if comparison['all_within_5pct'] else '部分未通过'}")
log(f" 输出目录: {output_dir}")
log(f" 日志文件: {log_path}")
log(f"{'='*70}")
log_file.close()
def _save_metrics(path, metrics):
"""保存评估指标到 JSON。"""
with open(path, "w") as f:
json.dump(metrics, f, indent=2, default=_json_default)
def _json_default(obj):
"""JSON 序列化的默认处理器。"""
if isinstance(obj, (np.integer,)):
return int(obj)
if isinstance(obj, (np.floating,)):
return float(obj)
if isinstance(obj, (np.ndarray,)):
return obj.tolist()
return str(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) | |
Collaborator
There was a problem hiding this comment.
权重放到model registry里，支持自动下载
|
|
||
| --- | ||
|
|
||
| ## Results |
Collaborator
There was a problem hiding this comment.
sample使用和train相同的config 合并
Collaborator
There was a problem hiding this comment.
不符合已有dataset格式，另外为什么要新建asu单独的dataset
Collaborator
There was a problem hiding this comment.
resources下的文件先建议都放到ppmat/utils/vocabs/crystals/下面
Collaborator
There was a problem hiding this comment.
构建structure对象在build里处理，
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实现
learncat163
force-pushed
the
feature/add-sgequidiff
branch
from
June 29, 2026 08:58
1f37fb0 to
f1ed062
Compare
Author
|
放弃了之前的 1:1的代码转译的paconvert的代码模式； 重新组织和复用了代码。 |
|
Thanks for your contribution! |
Collaborator
|
存在冲突 |
leeleolay
reviewed
Jul 16, 2026
leeleolay
left a comment
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的逻辑
Collaborator
There was a problem hiding this comment.
build structure 在数据集里处理
Collaborator
There was a problem hiding this comment.
构建structure对象在build里处理，
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 free
to 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
已经解决的有评审意见