- 
- 
- 
+
+
+
-[](https://www.star-history.com/#PaddlePaddle/PaddleMaterilas&type=date&legend=top-left) - Join the PaddleMaterials WeChat group to discuss with us!
-## Contribute to PaddleMaterials +## 🛠️ Contribute to PaddleMaterials For developer, please refer to [architecture](docs/ARCHITECTURE_ch.md). @@ -193,7 +193,7 @@ PaddleMaterials is licensed under the [Apache License 2.0](LICENSE). --- -## Acknowledgements +## 🙏 Acknowledgements This repository references code from the following projects: diff --git a/ppmat/metrics/__init__.py b/ppmat/metrics/__init__.py index acb6d511..1f6b904b 100644 --- a/ppmat/metrics/__init__.py +++ b/ppmat/metrics/__init__.py @@ -19,6 +19,7 @@ from ppmat.metrics.csp_metric import CSPMetric from ppmat.metrics.diffnmr_streaming_adapter import DiffNMRStreamingAdapter from ppmat.metrics.sfin_metric import SFINStreamingAdapter +from ppmat.metrics.sun_metric_utils import SUNMetric __all__ = [ "build_metric", @@ -27,6 +28,7 @@ "SFINStreamingAdapter", # "DiffNMRMetric", # "NLL", "CrossEntropyMetric", "SumExceptBatchMetric", "SumExceptBatchKL", + "SUNMetric", ] diff --git a/ppmat/metrics/sun_metric_utils.py b/ppmat/metrics/sun_metric_utils.py new file mode 100644 index 00000000..615006b7 --- /dev/null +++ b/ppmat/metrics/sun_metric_utils.py @@ -0,0 +1,294 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os +from collections import defaultdict +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +import numpy as np +import pandas as pd +from pymatgen.analysis.structure_matcher import StructureMatcher +from pymatgen.core import Structure +from tqdm import tqdm + +from ppmat.datasets.build_structure import BuildStructure + +logger = logging.getLogger(__name__) + +_chgnet_model = None +_graph_converter = None + + +def _load_chgnet(model_name: str = "chgnet_mptrj"): + global _chgnet_model, _graph_converter + if _chgnet_model is not None: + return _chgnet_model, _graph_converter + from ppmat.models import build_model_from_name + from ppmat.models.chgnet.chgnet_graph_converter import CHGNetGraphConverter + _graph_converter = CHGNetGraphConverter(atom_graph_cutoff=6.0, bond_graph_cutoff=3.0) + model, _ = build_model_from_name(model_name) + model.eval() + _chgnet_model = model + return _chgnet_model, _graph_converter + + +def predict_energy( + structures: List[Optional[Structure]], + model_name: str = "chgnet_mptrj", +) -> List[Optional[float]]: + model, converter = _load_chgnet(model_name) + energies = [] + for struct in tqdm(structures, desc="Predicting energies (CHGNet)"): + if struct is None: + energies.append(None) + continue + try: + graph = converter(struct) + pred = model.predict(graph) + e_per_atom = float(np.array(pred["energy_per_atom"]).flatten()[0]) + energies.append(e_per_atom) + except Exception as e: + logger.debug(f"Energy prediction failed: {e}") + energies.append(None) + return energies + + +def _composition_hash(structure: Structure) -> str: + return str(sorted(list(structure.atomic_numbers))) + + +def _parse_structures(raw_list, format_str): + results = [] + for item in raw_list: + try: + results.append(BuildStructure.build_one(item, format_str, niggli=False, canocial=False)) + except Exception: + results.append(None) + return results + + +def _match_against( + structures: List[Optional[Structure]], + reference_by_comp: dict, + matcher: StructureMatcher, + symmetric: bool = True, + dynamic: bool = False, +) -> List[bool]: + results = [] + for struct in tqdm(structures, desc="Matching structures"): + if struct is None: + results.append(False) + continue + h = _composition_hash(struct) + refs = reference_by_comp.get(h) + if refs is None: + results.append(True) + else: + for ref in refs: + if matcher.fit(struct, ref, symmetric=symmetric): + results.append(False) + break + else: + results.append(True) + if dynamic and results[-1]: + reference_by_comp[h].append(struct) + return results + + +def compute_uniqueness( + structures: List[Optional[Structure]], + stol: float = 0.5, + angle_tol: float = 10.0, + ltol: float = 0.3, + attempt_supercell: bool = True, + symmetric: bool = True, +) -> List[bool]: + matcher = StructureMatcher(stol=stol, angle_tol=angle_tol, ltol=ltol, attempt_supercell=attempt_supercell) + return _match_against(structures, defaultdict(list), matcher, symmetric=symmetric, dynamic=True) + + +def compute_novelty( + structures: List[Optional[Structure]], + reference_structures: List[Structure], + stol: float = 0.5, + angle_tol: float = 10.0, + ltol: float = 0.3, + attempt_supercell: bool = True, + symmetric: bool = True, +) -> List[bool]: + matcher = StructureMatcher(stol=stol, angle_tol=angle_tol, ltol=ltol, attempt_supercell=attempt_supercell) + reference_by_comp = defaultdict(list) + for ref in reference_structures: + reference_by_comp[_composition_hash(ref)].append(ref) + return _match_against(structures, reference_by_comp, matcher, symmetric=symmetric) + + +def compute_stability( + energy_above_hull: List[Optional[float]], + threshold: float = 0.0, +) -> List[bool]: + return [ + (eah is not None and not np.isnan(eah) and eah < threshold) + for eah in energy_above_hull + ] + + +def compute_sun( + stability: List[bool], + uniqueness: List[bool], + novelty: List[bool], + structures: List[Optional[Structure]], + min_elements: int = 2, +) -> Dict[str, float]: + n = len(stability) + non_trivial = [] + for s in structures: + if s is not None and len(set(s.composition)) >= min_elements: + non_trivial.append(True) + else: + non_trivial.append(False) + + sun = [s and u and nv and nt for s, u, nv, nt in zip(stability, uniqueness, novelty, non_trivial)] + + def rate(flags): + return round(100 * sum(flags) / max(len(flags), 1), 2) + + return { + "total": n, + "valid": sum(1 for s in structures if s is not None), + "non_trivial": sum(non_trivial), + "stability_rate": rate(stability), + "uniqueness_rate": rate(uniqueness), + "novelty_rate": rate(novelty), + "sun_rate": rate(sun), + "sun_count": sum(sun), + } + + +class SUNMetric: + def __init__( + self, + gt_file_path: Optional[str] = None, + stol: float = 0.5, + angle_tol: float = 10.0, + ltol: float = 0.3, + stability_threshold: float = 0.0, + attempt_supercell: bool = True, + use_chgnet: bool = False, + ): + self.gt_file_path = gt_file_path + self.stol = stol + self.angle_tol = angle_tol + self.ltol = ltol + self.stability_threshold = stability_threshold + self.attempt_supercell = attempt_supercell + self.use_chgnet = use_chgnet + self._reference_structures = None + + def _load_reference(self): + if self._reference_structures is not None: + return + if self.gt_file_path is None: + return + if not os.path.exists(self.gt_file_path): + logger.warning(f"Reference file not found: {self.gt_file_path}") + return + df = pd.read_csv(self.gt_file_path) + if "cif" not in df.columns: + logger.warning("Reference CSV must have a 'cif' column") + return + self._reference_structures = [] + for cif_str in tqdm(df["cif"], desc="Loading reference structures"): + try: + self._reference_structures.append(BuildStructure.build_one(cif_str, "cif_str", niggli=False, canocial=False)) + except Exception: + pass + logger.info(f"Loaded {len(self._reference_structures)} reference structures") + + def __call__( + self, + generated: Union[List[dict], List[str], pd.DataFrame], + energy_above_hull: Optional[List[Optional[float]]] = None, + ) -> Dict[str, float]: + if isinstance(generated, pd.DataFrame): + if "cif" in generated.columns: + generated = generated["cif"].tolist() + elif "structure" in generated.columns: + generated = generated["structure"].tolist() + else: + raise ValueError("DataFrame must have 'cif' or 'structure' column") + + if isinstance(generated, list) and len(generated) > 0: + fmt = None + if isinstance(generated[0], str): + fmt = "cif_str" + elif isinstance(generated[0], dict): + fmt = "array" + structures = _parse_structures(generated, fmt) if fmt else generated + else: + raise ValueError("generated must be a non-empty list or DataFrame") + + n = len(structures) + logger.info(f"Evaluating {n} generated structures") + + uniqueness = compute_uniqueness( + structures, stol=self.stol, angle_tol=self.angle_tol, ltol=self.ltol, + attempt_supercell=self.attempt_supercell, + ) + + novelty_list = [True] * n + self._load_reference() + if self._reference_structures: + novelty_list = compute_novelty( + structures, self._reference_structures, + stol=self.stol, angle_tol=self.angle_tol, ltol=self.ltol, + attempt_supercell=self.attempt_supercell, + ) + + if energy_above_hull is not None: + stability = compute_stability(energy_above_hull, threshold=self.stability_threshold) + energies = None + elif self.use_chgnet: + energies = predict_energy(structures) + stability = [(e is not None and e < self.stability_threshold) for e in energies] + else: + stability = [False] * n + energies = None + logger.warning("No energy_above_hull and use_chgnet=False, stability will be 0%") + + results = compute_sun(stability, uniqueness, novelty_list, structures) + + if energies is not None: + valid = [e for e in energies if e is not None] + if valid: + results["avg_energy_per_atom"] = float(np.mean(valid)) + + if energy_above_hull is not None: + meta = compute_stability(energy_above_hull, threshold=0.1) + msun = compute_sun(meta, uniqueness, novelty_list, structures) + results["metastable_rate"] = msun["stability_rate"] + results["msun_rate"] = msun["sun_rate"] + results["msun_count"] = msun["sun_count"] + + logger.info( + f"S.U.N.: Stability={results['stability_rate']:.2f}%, " + f"Uniqueness={results['uniqueness_rate']:.2f}%, " + f"Novelty={results['novelty_rate']:.2f}%, " + f"S.U.N.={results['sun_rate']:.2f}%" + ) + return results diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py index d258bd2a..f5013620 100644 --- a/ppmat/models/__init__.py +++ b/ppmat/models/__init__.py @@ -42,6 +42,7 @@ from ppmat.models.megnet.megnet import MEGNetPlus from ppmat.models.infgcn.infgcn import InfGCN from ppmat.models.mateno.mateno import MatENO +from ppmat.models.miad.miad import MiAD from ppmat.models.sfin.sfin import SFIN from ppmat.utils import download from ppmat.utils import logger @@ -68,6 +69,7 @@ "DiffNMR", "InfGCN", "MatENO", + "MiAD", "SFIN", ] @@ -113,6 +115,7 @@ "mattergen_ml2ddb": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb.zip", "mattergen_ml2ddb_chemical_system": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_chemical_system.zip", "mattergen_ml2ddb_space_group": "https://paddle-org.bj.bcebos.com/paddlematerial/workflow/ml2ddb/mattergen_ml2ddb_space_group.zip", + "miad_mp20": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/MiAD/miad_mp20.zip", "sfin_haadf_enhance": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_haadf_enhance.zip", "sfin_haadf_detect": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_haadf_detect.zip", "sfin_bf_enhance": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_bf_enhance.zip", diff --git a/ppmat/models/miad/__init__.py b/ppmat/models/miad/__init__.py new file mode 100644 index 00000000..2a5bb807 --- /dev/null +++ b/ppmat/models/miad/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ppmat.models.miad.miad import MiAD + +__all__ = [ + "MiAD", +] diff --git a/ppmat/models/miad/crystal_diffusion.py b/ppmat/models/miad/crystal_diffusion.py new file mode 100644 index 00000000..319628d0 --- /dev/null +++ b/ppmat/models/miad/crystal_diffusion.py @@ -0,0 +1,364 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import paddle +from paddle_scatter import scatter_mean + +from ppmat.models.common.time_embedding import SinusoidalTimeEmbeddings +from ppmat.models.miad.type_diffusion import D3PM +from ppmat.models.miad.type_diffusion import DDPMOnehot +from ppmat.schedulers import build_scheduler +from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal +from ppmat.utils.crystal import lattice_params_to_matrix_paddle + +_DEFAULT_GAMMA = { + "csp_perov5": 5e-7, "gen_perov5": 5e-7, + "csp_mp20": 1e-5, "gen_mp20": 1e-5, + "csp_alex_mp20": 1e-5, "gen_alex_mp20": 1e-5, + "csp_mpts52": 1e-5, "gen_mpts52": 1e-5, + "csp_carbon24": 5e-7, "gen_carbon24": 1e-5, +} + + +def parse_num_atoms_to_per_crystal(num_atoms_data): + if num_atoms_data is None: + return None + if hasattr(num_atoms_data, "numpy"): + num_atoms_np = num_atoms_data.numpy().flatten() + elif hasattr(num_atoms_data, "reshape"): + num_atoms_np = num_atoms_data.reshape(-1) + else: + num_atoms_np = np.array(num_atoms_data).flatten() + return paddle.to_tensor(num_atoms_np.astype("int64")), num_atoms_np + + +def mean_interleave(t, num_repeats): + batch_idx = paddle.repeat_interleave( + paddle.arange(num_repeats.shape[0]), num_repeats + ) + return scatter_mean(t, batch_idx, dim=0) + + +class CrystalGen: + """Crystal generation orchestrator. + + Directly uses ppmat schedulers (DDPMScheduler, ScoreSdeVeSchedulerWrapped) + for coefficient pre-computation, with custom step logic for each diffusion + type. + """ + + def __init__(self, diffusion_config, logger): + self.config = diffusion_config + self.logger = logger + self.cont_time = self.config["cont_time"] + self.num_steps = self.config["num_steps"] + self.eps = 1e-3 + self.time_embedding = SinusoidalTimeEmbeddings( + self.config.get("time_embed_dim", 256) + ) + self.gen_type = "gen" in self.config["task"] + + # Lattice diffusion + lat_cfg = self.config["lat_diffusion"] + self.lat_method = lat_cfg["method"] + self.lat_scheduler = build_scheduler(lat_cfg.get("scheduler_cfg", { + "__class_name__": "DDPMScheduler", + "__init_params__": { + "num_train_timesteps": self.num_steps, + "beta_schedule": "squaredcos_cap_v2", + }, + })) + if self.lat_method == "fm_lenang": + self._lenang2lat = None + self.gamma_alpha, self.gamma_theta = 1.3, 0.25 + self.angle_difference_bound = 20 + + # Frac diffusion + frac_cfg = self.config["frac_diffusion"] + self.frac_method = frac_cfg["method"] + self.frac_scheduler = build_scheduler(frac_cfg.get("scheduler_cfg", { + "__class_name__": "ScoreSdeVeSchedulerWrapped", + "__init_params__": { + "num_train_timesteps": self.num_steps, + "sigma_min": 0.005, + "sigma_max": 0.5, + "sampling_eps": 1e-3, + }, + })) + self.step_lr = frac_cfg.get("step_lr", None) + if self.step_lr is None: + self.step_lr = _DEFAULT_GAMMA.get(self.config["task"], 1e-5) + self.sigmas_t = self.frac_scheduler.discrete_sigmas[:, None] + self.sigmas_norm_t = self.frac_scheduler.discrete_sigmas_norm[:, None] + self.sb = 0.005 + + # Type diffusion + if self.gen_type: + type_cfg = self.config["type_diffusion"] + switch_type = { + "ddpm_onehot": DDPMOnehot, + "d3pm": D3PM, + } + self.type_diffusion = switch_type[type_cfg["method"]](self.config) + else: + self.type_diffusion = None + + def _time_sample(self, batch): + t = paddle.rand([batch["batch_size"]]) + t = self.eps + (self.num_steps - 1 - self.eps) * t + if not self.cont_time: + t = t.round().cast("int64") + t_per_atom = t.repeat_interleave(batch["num_atoms"]) + return [t, t_per_atom] + + def _time_iterator(self, batch, start_from=-1): + if start_from == -1: + start_from = self.num_steps - 1 + for t_val in range(start_from, -1, -1): + t = paddle.full([batch["batch_size"]], t_val, dtype="float32") + yield [t, t.repeat_interleave(batch["num_atoms"])] + + def _lenang2lat_fn(self, la): + if self._lenang2lat is None: + self._lenang2lat = lambda la: lattice_params_to_matrix_paddle( + la[:, :3], la[:, 3:] + ) + return self._lenang2lat(la) + + def forward_step_sample(self, x0, t, batch): + l0, f0, a0 = x0 + t0_idx = t[0].cast("int64") + t1_idx = t[1].cast("int64") + + # Lattice forward + if self.lat_method == "ddpm": + noise = paddle.randn(l0.shape) + self.lat_randn = noise + lt = self.lat_scheduler.add_noise(l0, noise, t0_idx) + elif self.lat_method == "fm": + noise = paddle.randn(l0.shape) + step = (1 + t[0][:, None, None]) / self.num_steps + lt = (1 - step) * l0 + step * noise + self.lat_ut = -noise + elif self.lat_method == "fm_lenang": + xT = self._prior_lat(batch) + step = (1 + t[0][:, None, None]) / self.num_steps + lt = (1 - step) * l0 + step * xT + self.lat_ut = l0 - xT + else: + lt = l0 + + # Frac forward + if self.frac_method == "wrapped_normal": + noise = paddle.randn(f0.shape) + self.frac_randn = noise + ft = (f0 + self.sigmas_t[t1_idx] * noise) % 1.0 + elif self.frac_method == "pfm": + xT_minus_x0 = paddle.rand(f0.shape) - 0.5 + step = (1 + t[1][:, None]) / self.num_steps + ft = (f0 + step * xT_minus_x0) % 1.0 + self.frac_ut = -xT_minus_x0 + else: + ft = f0 + + # Type forward + if self.gen_type: + at = self.type_diffusion.forward_step_sample(a0, t[1], batch) + else: + at = a0 + + return [lt, ft, at] + + def reverse_step_sample(self, xt, t, model, batch): + lt, ft, at = xt + + if self.config["method"] == "DiffCSP": + _, f_pred, _ = self.model_prediction(xt, t, model, batch) + ft_05 = self.frac_reverse_part1(f_pred, ft, t[1]) + xt_05 = [lt, ft_05, at] + l_pred, f_pred, a_pred = self.model_prediction(xt_05, t, model, batch) + lt_1 = self.lat_reverse(l_pred, lt, t[0]) + ft_1 = self.frac_reverse_part2(f_pred, ft_05, t[1]) + else: + batch["prediction"] = self.model_prediction(xt, t, model, batch) + l_pred, f_pred, a_pred = batch["prediction"] + lt_1 = self.lat_reverse(l_pred, lt, t[0]) + ft_1 = self.frac_reverse(f_pred, ft, t[1]) + + at_1 = ( + self.type_diffusion.reverse_step_sample(a_pred, at, t[1], batch) + if self.gen_type else at + ) + return [lt_1, ft_1, at_1] + + def lat_reverse(self, pred, xt, t): + t_idx = t.cast("int64") + if self.lat_method == "ddpm": + return self.lat_scheduler.step(pred, t_idx[0], xt).prev_sample + elif self.lat_method == "fm": + pred = -pred + step = (1 + t[:, None, None]) / self.num_steps + x0_pred = (xt - step * pred) / (1 - step) + vt = x0_pred - pred + return xt + step / self.num_steps * vt + elif self.lat_method == "fm_lenang": + vt = pred + return xt + vt / self.num_steps + return xt + + def frac_reverse(self, pred, xt, t): + t_idx = t.cast("int64") + if self.frac_method == "wrapped_normal": + return self.frac_reverse_part2(pred, xt, t) + elif self.frac_method == "pfm": + return (xt + pred / self.num_steps) % 1.0 + return xt + + def frac_reverse_part1(self, pred, xt, t): + t_idx = t.cast("int64") + st = self.sigmas_t[t_idx] + snt = self.sigmas_norm_t[t_idx] + step_size = self.step_lr * (st / self.sb) ** 2 + drift = -step_size * pred * paddle.sqrt(snt) + diffusion = paddle.sqrt(2 * step_size) * paddle.randn(xt.shape) + return xt + drift + diffusion + + def frac_reverse_part2(self, pred, xt, t): + t_idx = t.cast("int64") + st = self.sigmas_t[t_idx] + st_1 = self.sigmas_t[paddle.maximum(t_idx - 1, paddle.to_tensor(0))] + snt = self.sigmas_norm_t[t_idx] + step_size = st ** 2 - st_1 ** 2 + drift = -step_size * pred * paddle.sqrt(snt) + diffusion = paddle.sqrt(st_1 ** 2 * (st ** 2 - st_1 ** 2) / (st ** 2)) * paddle.randn(xt.shape) + return (xt + drift + diffusion) % 1.0 + + def _prior_lat(self, batch): + if self.lat_method == "ddpm": + return paddle.randn([batch["batch_size"], 3, 3], dtype="float32") + elif self.lat_method == "fm": + return paddle.randn([batch["batch_size"], 3, 3], dtype="float32") + elif self.lat_method == "fm_lenang": + bs = batch["batch_size"] + la = paddle.zeros([bs, 6], dtype="float32") + gamma = paddle.distribution.Gamma( + paddle.to_tensor([self.gamma_alpha], dtype="float32"), + paddle.to_tensor([self.gamma_theta], dtype="float32"), + ) + la[:, :3] = 2 + gamma.sample([bs, 3]) + collected = [] + while len(collected) < bs: + ang = 60 + 60 * paddle.rand([bs, 3]) + check = ( + (ang[:, 0] + ang[:, 1] - ang[:, 2] > self.angle_difference_bound) + * (ang[:, 2] + ang[:, 0] - ang[:, 1] > self.angle_difference_bound) + * (ang[:, 1] + ang[:, 2] - ang[:, 0] > self.angle_difference_bound) + ) + collected.append(ang[check]) + la[:, 3:] = paddle.concat(collected)[:bs] + return self._lenang2lat_fn(la) + raise ValueError(f"Unknown lat_method: {self.lat_method}") + + def _prior_frac(self, batch): + na = batch["num_atoms"] + total = int(na.sum()) if na.ndim > 0 else int(na) + if self.frac_method == "wrapped_normal": + return paddle.rand([total, 3], dtype="float32") + elif self.frac_method == "pfm": + return paddle.rand([total, 3]) + raise ValueError(f"Unknown frac_method: {self.frac_method}") + + def prior_sample(self, batch): + prior_at = ( + self.type_diffusion.prior_sample(batch) + if self.gen_type else batch["atom_types"] + ) + return [self._prior_lat(batch), self._prior_frac(batch), prior_at] + + def model_prediction(self, xt, t, model, batch): + lt, ft, at = xt + time_emb = self.time_embedding(1000 * (t[0] / self.num_steps) + 1) + nn_pred = model(time_emb, at, ft, lt, batch["num_atoms"], batch["batch_idx"]) + return [ + nn_pred[0], + nn_pred[1], + nn_pred[2] if self.gen_type else None, + ] + + def train_step(self, batch, model, mode): + batch["t"] = self._time_sample(batch) + batch["xt"] = self.forward_step_sample(batch["x0"], batch["t"], batch) + batch["prediction"] = self.model_prediction(batch["xt"], batch["t"], model, batch) + + # Lattice loss + if self.lat_method == "ddpm": + loss_lat = ((batch["prediction"][0] - self.lat_randn) ** 2).reshape([-1, 9]).mean(axis=1).mean() + elif self.lat_method in ("fm", "fm_lenang"): + loss_lat = ((batch["prediction"][0] - self.lat_ut) ** 2).reshape([-1, 9]).mean(axis=1).mean() + else: + loss_lat = paddle.to_tensor(0.0) + + # Frac loss + if self.frac_method == "wrapped_normal": + t_idx = batch["t"][1].cast("int64") + st = self.sigmas_t[t_idx] + snt = self.sigmas_norm_t[t_idx] + normed_score = d_log_p_wrapped_normal(st * self.frac_randn, st) / paddle.sqrt(snt) + loss_frac = ((batch["prediction"][1] - normed_score) ** 2).reshape([-1, 3]).mean(axis=1).mean() + elif self.frac_method == "pfm": + loss_frac = ((batch["prediction"][1] - self.frac_ut) ** 2).reshape([-1, 3]).mean(axis=1).mean() * 10 + else: + loss_frac = paddle.to_tensor(0.0) + + batch["loss"] = loss_lat + loss_frac + + if self.gen_type: + loss_type = self.type_diffusion.loss(batch).mean() + batch["loss"] = batch["loss"] + loss_type + + if self.logger is not None: + t_log = paddle.clip(batch["t"][0].clone().cast("int64"), 0, 999) + self.logger.add(f"loss:lattice:{mode}", loss_lat.item(), stack_after_epoch=True) + self.logger.add(f"loss:coord:{mode}", loss_frac.item(), stack_after_epoch=True) + loss_lat_t = paddle.zeros([self.num_steps], dtype=loss_lat.dtype) + loss_lat_t[t_log] = batch["prediction"][0].reshape([-1, 9]).mean(axis=1) + self.logger.add(f"loss:lattice4time:{mode}", loss_lat_t, stack_after_epoch=True) + loss_frac_t = paddle.zeros([self.num_steps], dtype=loss_frac.dtype) + loss_frac_t[t_log] = mean_interleave( + batch["prediction"][1].reshape([-1, 3]).mean(axis=1), batch["num_atoms"], + ) + self.logger.add(f"loss:coord4time:{mode}", loss_frac_t, stack_after_epoch=True) + if self.gen_type: + self.logger.add(f"loss:type:{mode}", loss_type.item(), stack_after_epoch=True) + + return batch + + def sampling_procedure(self, model, batch, progress_printer): + batch["xt"] = self.prior_sample(batch) + for t_vec in self._time_iterator(batch, start_from=self.num_steps - 1): + batch["t"] = t_vec + progress_printer(batch["t"][0][0].item()) + batch["xt"] = self.reverse_step_sample(batch["xt"], batch["t"], model, batch) + batch["xt"] = self.output_transform(batch["xt"], batch) + batch["x0_prediction"] = batch["xt"] + return batch + + def output_transform(self, x0, batch): + _out = lambda diff, x: diff.output_transform(x, batch) if hasattr(diff, 'output_transform') else x + return [ + x0[0], + x0[1], + _out(self.type_diffusion, x0[2]) if self.gen_type else x0[2], + ] diff --git a/ppmat/models/miad/miad.py b/ppmat/models/miad/miad.py new file mode 100644 index 00000000..b7d620d4 --- /dev/null +++ b/ppmat/models/miad/miad.py @@ -0,0 +1,225 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import paddle +import paddle.nn as nn + +from ppmat.models.miad.crystal_diffusion import CrystalGen +from ppmat.models.miad.crystal_diffusion import parse_num_atoms_to_per_crystal +from ppmat.models.diffcsp.diffcsp import CSPNet +from ppmat.utils.crystal import lattices_to_params_shape_numpy + + +class MiADCSPNet(CSPNet): + """CSPNet with block-diagonal edge generation. + Uses block_diag instead of meshgrid to avoid GPU crash on certain batch sizes. + """ + + def gen_edges(self, num_atoms, frac_coords): + lis = [paddle.ones([int(n), int(n)], dtype="int64") for n in num_atoms] + fc_graph = paddle.block_diag(lis) + fc_edges = paddle.nonzero(fc_graph).t() + return fc_edges, frac_coords[fc_edges[1]] - frac_coords[fc_edges[0]] + + +def _to_numpy(x): + if hasattr(x, "numpy"): + return x.numpy() + return np.asarray(x) + + +def _extract_x0(batch): + if "x0" in batch: + return batch + sa = batch.get("structure_array", batch) + num_atoms_np = _to_numpy(sa["num_atoms"]).flatten().astype("int64") + frac_coords_np = _to_numpy(sa["frac_coords"]).astype("float32") + atom_types_np = _to_numpy(sa["atom_types"]).flatten().astype("int64") + lattice_np = _to_numpy(sa["lattice"]).astype("float32") + if lattice_np.ndim == 3 and lattice_np.shape[0] == 1: + lattice_np = lattice_np.reshape(3, 3) + batch_size = len(num_atoms_np) + total_atoms = int(num_atoms_np.sum()) + batch_idx_np = np.concatenate( + [np.full(int(n), i) for i, n in enumerate(num_atoms_np)] + ).astype("int64") + batch["x0"] = [ + paddle.to_tensor(lattice_np.reshape(batch_size, 3, 3)), + paddle.to_tensor(frac_coords_np.reshape(total_atoms, 3)), + paddle.to_tensor(atom_types_np.reshape(total_atoms)), + ] + batch["num_atoms"] = paddle.to_tensor(num_atoms_np) + batch["batch_idx"] = paddle.to_tensor(batch_idx_np) + batch["atom_types"] = paddle.to_tensor(atom_types_np) + batch["batch_size"] = batch_size + return batch + + +class MiAD(nn.Layer): + """Mirage Atom Diffusion model.""" + + def __init__(self, model_cfg=None, diffusion_cfg=None, **kwargs): + super().__init__() + + model_cfg = model_cfg or {} + diffusion_cfg = diffusion_cfg or {} + + model_cfg = dict(model_cfg) + model_cfg.setdefault("num_classes", model_cfg.pop("max_atoms", 100)) + _cspnet_keys = { + "hidden_dim", "latent_dim", "num_layers", "act_fn", "dis_emb", + "num_freqs", "edge_style", "ln", "ip", "smooth", "pred_type", + "prop_dim", "pred_scalar", "num_classes", + } + cspnet_kwargs = {k: v for k, v in model_cfg.items() if k in _cspnet_keys} + self.decoder = MiADCSPNet(**cspnet_kwargs) + self.diffusion = CrystalGen(diffusion_cfg, logger=None) + + def set_state_dict(self, state_dict, use_structured_name=True): + if not any(k.startswith("decoder.") for k in state_dict.keys()): + state_dict = {f"decoder.{k}": v for k, v in state_dict.items()} + param_state = {} + for name, param in self.named_parameters(): + if name in state_dict: + v = state_dict[name] + if hasattr(v, "numpy"): + v = v.numpy() + elif not isinstance(v, np.ndarray): + v = np.asarray(v) + if name.endswith(".weight") and len(v.shape) == 2: + v = v.T + if v.shape == param.shape: + param_state[name] = v.astype(param.numpy().dtype) + for name in param_state: + self.get_parameter(name).set_value(param_state[name]) + return [], [] + + def forward(self, batch, **kwargs): + mode = "train" if self.training else "val" + batch = _extract_x0(batch) + batch = self.diffusion.train_step( + batch=batch, + model=self.decoder, + mode=mode, + ) + loss = batch["loss"] + loss_dict = { + "loss": loss, + } + return {"loss_dict": loss_dict} + + @paddle.no_grad() + def sample(self, batch_data, num_inference_steps=None): + if "structure_array" in batch_data: + num_atoms_data = batch_data["structure_array"].get("num_atoms", None) + else: + num_atoms_data = batch_data.get("num_atoms", None) + + if "batch_idx" in batch_data: + pass + elif num_atoms_data is not None: + parsed = parse_num_atoms_to_per_crystal(num_atoms_data) + if parsed is not None: + num_atoms, num_atoms_np = parsed + batch_size = len(num_atoms_np) + batch_idx_np = np.concatenate( + [np.full(int(n), i) for i, n in enumerate(num_atoms_np)] + ) + batch_idx = paddle.to_tensor(batch_idx_np.astype("int64")) + atom_types = paddle.zeros([int(num_atoms_np.sum())], dtype="int64") + batch_data = { + **batch_data, + "num_atoms": num_atoms, + "batch_idx": batch_idx, + "atom_types": atom_types, + "batch_size": batch_size, + } + + if num_inference_steps is not None: + original_steps = self.diffusion.num_steps + self.diffusion.num_steps = num_inference_steps + else: + original_steps = None + + def progress_printer(t): + return None + + try: + batch = self.diffusion.sampling_procedure( + model=self.decoder, + batch=batch_data, + progress_printer=progress_printer, + ) + finally: + if original_steps is not None: + self.diffusion.num_steps = original_steps + + # Extract results + x0_pred = batch["x0_prediction"] + lattices = x0_pred[0] + frac_coords = x0_pred[1] + atom_types = x0_pred[2] + + result = [] + num_atoms = batch_data.get("num_atoms", None) + batch_size = batch_data.get("batch_size", lattices.shape[0]) + + start_idx = 0 + for i in range(batch_size): + if num_atoms is not None: + n = int(num_atoms[i]) + else: + n = frac_coords.shape[0] + if i > 0: + break + # Convert tensors to numpy so downstream consumers (BuildStructure, + # CSPMetric) can handle them in multi-process p_map without + # serializing paddle tensors. + lat_i = lattices[i] + fc_i = frac_coords[start_idx : start_idx + n] + at_i = atom_types[start_idx : start_idx + n] + lat_np = _to_numpy(lat_i) + fc_np = _to_numpy(fc_i) + at_np = _to_numpy(at_i) + start_idx += n + # Filter out mirage atoms (atom_types == 0) produced by Mirage + # Infusion, mirroring lib/data/crystal_data_storage.py save_batch. + valid_mask = at_np != 0 + if valid_mask.any(): + at_np = at_np[valid_mask] + fc_np = fc_np[valid_mask] + n = int(valid_mask.sum()) + else: + # All atoms are mirage atoms; keep original count but replace + # invalid type 0 with 1 (H) to avoid empty structure errors. + at_np = np.where(at_np == 0, 1, at_np) + # Pre-compute lattice params so BuildStructure does not need to + # derive them from a (possibly non-list) lattice matrix. + lat_for_params = ( + lat_np.reshape(1, 3, 3) if lat_np.ndim == 2 else lat_np + ) + lengths, angles = lattices_to_params_shape_numpy(lat_for_params) + result.append( + { + "num_atoms": n, + "atom_types": at_np, + "frac_coords": fc_np, + "lattice": lat_np, + "lengths": lengths.flatten(), + "angles": angles.flatten(), + } + ) + + return {"result": result} diff --git a/ppmat/models/miad/type_diffusion.py b/ppmat/models/miad/type_diffusion.py new file mode 100644 index 00000000..37b83f3e --- /dev/null +++ b/ppmat/models/miad/type_diffusion.py @@ -0,0 +1,206 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import paddle +import paddle.nn.functional as F + +from ppmat.schedulers import build_scheduler + + +class DDPMOnehot: + """DDPM for atom types with one-hot encoding.""" + + def __init__(self, diffusion_config): + self.num_steps = diffusion_config["num_steps"] + self.num_types = 100 + self.scheduler = build_scheduler( + diffusion_config["type_diffusion"].get("scheduler_cfg", { + "__class_name__": "DDPMScheduler", + "__init_params__": { + "num_train_timesteps": self.num_steps, + "beta_schedule": "squaredcos_cap_v2", + }, + }) + ) + self.to_domain = lambda types: F.one_hot(types - 1, num_classes=self.num_types).cast("float32") + self.from_domain = lambda onehot: onehot.argmax(axis=-1) + 1 + + def forward_step_sample(self, x0, t, batch): + onehot_x0 = self.to_domain(x0) + noise = paddle.randn(onehot_x0.shape) + self.randn_x = noise + t_idx = t.cast("int64") + return self.scheduler.add_noise(onehot_x0, noise, t_idx) + + def reverse_step_sample(self, eps_pred, onehot_xt, t, batch): + t_idx = t.cast("int64") + out = self.scheduler.step(eps_pred, t_idx[0], onehot_xt) + onehot_xt_1 = out.prev_sample + if (t_idx == 0).all(): + return self.from_domain(onehot_xt_1) + return onehot_xt_1 + + def prior_sample(self, batch): + na = batch["num_atoms"] + total = int(na.sum()) if na.ndim > 0 else int(na) + return paddle.randn([total, self.num_types], dtype="float32") + + def loss(self, batch): + eps_pred = batch["prediction"][2] + return ((eps_pred - self.randn_x) ** 2).reshape([eps_pred.shape[0], -1]).mean(axis=1) + + +class D3PMUniformScheduler: + """D3PM scheduler with uniform transition matrices (model-internal). + + Pre-computes Q_t, cumprod_Q_t, Q_{t-1}, cumprod_Q_{t-1} for all timesteps + using a cosine schedule. This is the uniform-transition variant used by MiAD, + as opposed to the absorbing-state D3PMScheduler in ppmat.schedulers. + """ + + def __init__( + self, + num_train_timesteps: int = 1000, + num_types: int = 100, + s: float = 0.008, + ): + self.num_types = num_types + + discretization = paddle.arange( + 1, num_train_timesteps + 1, dtype="float64" + ) + f_t = paddle.cos( + (discretization / (num_train_timesteps + 1) + s) / (1 + s) * math.pi / 2 + ) + f_0 = paddle.cos( + (paddle.to_tensor(0.0, dtype="float64") + s) / (1 + s) * math.pi / 2 + ) + a_t = f_t / f_0 + cumprod_alphas_t = a_t + cumprod_alphas_t_1 = paddle.concat( + [paddle.to_tensor([1.0], dtype="float64"), cumprod_alphas_t[:-1]] + ) + betas_t = 1 - cumprod_alphas_t / cumprod_alphas_t_1 + + Q_t_list = [] + for t_idx in range(num_train_timesteps): + mat = paddle.full( + (num_types, num_types), betas_t[t_idx] / float(num_types) + ) + diag_val = 1 - betas_t[t_idx] * (num_types - 1) / num_types + for i in range(num_types): + mat[i, i] = diag_val.cast("float32") + Q_t_list.append(mat) + Q_t = paddle.stack(Q_t_list, axis=0) + + cumprod_Q_t_list = [Q_t[0]] + for t_idx in range(1, num_train_timesteps): + cumprod_Q_t_list.append( + paddle.matmul(cumprod_Q_t_list[-1], Q_t[t_idx]) + ) + cumprod_Q_t = paddle.stack(cumprod_Q_t_list, axis=0) + + Q_t_1 = paddle.concat( + [paddle.eye(num_types, dtype="float32").unsqueeze(0), Q_t[:-1]], + axis=0, + ) + + cumprod_Q_t_1 = paddle.concat( + [ + paddle.eye(num_types, dtype="float32").unsqueeze(0), + cumprod_Q_t[:-1], + ], + axis=0, + ) + + self.Q_t = Q_t.reshape([-1, num_types, num_types]) + self.Q_t_1 = Q_t_1.reshape([-1, num_types, num_types]) + self.cumprod_Q_t = cumprod_Q_t.reshape([-1, num_types, num_types]) + self.cumprod_Q_t_1 = cumprod_Q_t_1.reshape([-1, num_types, num_types]) + + +class D3PM: + """Discrete Denoising Diffusion Probabilistic Model for atom types.""" + + def __init__(self, diffusion_config): + self.config = diffusion_config["type_diffusion"] + self.num_steps = diffusion_config["num_steps"] + self.scheduler = D3PMUniformScheduler( + num_train_timesteps=self.num_steps, + num_types=100, + ) + self.num_types = self.scheduler.num_types + self.Q_t = self.scheduler.Q_t + self.Q_t_1 = self.scheduler.Q_t_1 + self.cumprod_Q_t = self.scheduler.cumprod_Q_t + self.cumprod_Q_t_1 = self.scheduler.cumprod_Q_t_1 + + self.to_domain = lambda types: F.one_hot(types, num_classes=self.num_types).cast("float32") + self.from_domain = lambda onehot: onehot.argmax(axis=-1) + self.prediction_to_domain = lambda pred: F.softmax(pred, axis=-1) + self.default_loss_scale = 1000 + + def output_transform(self, x0, batch): + return self.from_domain(x0) + + def forward_step_sample(self, x0, t, batch): + onehot_x0 = self.to_domain(x0) + xt_probs = paddle.matmul(onehot_x0[:, None, :], self.cumprod_Q_t[t.cast("int64")])[:, 0, :] + xt = paddle.distribution.Categorical(logits=paddle.log(xt_probs.clip(1e-12))).sample().cast("int64") + return self.to_domain(xt) + + def _reverse_step_distribution(self, onehot_x0, onehot_xt, t): + t_idx = t.cast("int64") + numerator = ( + paddle.matmul(onehot_xt[:, None, :], self.Q_t[t_idx].transpose([0, 2, 1]))[:, 0, :] + * paddle.matmul(onehot_x0[:, None, :], self.cumprod_Q_t_1[t_idx])[:, 0, :] + ) + denominator = ( + paddle.matmul(onehot_x0[:, None, :], self.cumprod_Q_t[t_idx])[:, 0, :] * onehot_xt + ).sum(axis=-1)[:, None] + result = numerator / (denominator + 1e-8) + result = result / result.sum(axis=-1, keepdim=True) + return paddle.where(paddle.isnan(result), paddle.full_like(result, 1.0 / self.num_types), result) + + def reverse_step_sample(self, onehot_pred, onehot_xt, t, batch): + onehot_x0 = self.prediction_to_domain(onehot_pred) + xt_1_probs = self._reverse_step_distribution(onehot_x0, onehot_xt, t) + if (t.cast("int64") == 0).all(): + return self.to_domain(self.from_domain(xt_1_probs.cast("float32"))) + xt_1 = paddle.distribution.Categorical(logits=paddle.log(xt_1_probs.clip(1e-12))).sample().cast("int64") + return self.to_domain(xt_1) + + def prior_sample(self, batch): + na = batch["num_atoms"] + total = int(na.sum()) if na.ndim > 0 else int(na) + shape = [total, self.num_types] + xT_probs = paddle.ones(shape, dtype="float32") / self.num_types + xT = paddle.distribution.Categorical(logits=paddle.log(xT_probs.clip(1e-12))).sample().cast("int64") + return self.to_domain(xT) + + def loss(self, batch): + onehot_xt = batch["xt"][2] + t = batch["t"][1].cast("int64") + onehot_x0_pred = self.prediction_to_domain(batch["prediction"][2]) + pred_xt_1_probs = self._reverse_step_distribution(onehot_x0_pred, onehot_xt, t) + onehot_x0 = self.to_domain(batch["x0"][2]) + orig_xt_1_probs = self._reverse_step_distribution(onehot_x0, onehot_xt, t) + eps = 1e-8 + kl_loss = ( + (orig_xt_1_probs * (paddle.log(orig_xt_1_probs + eps) - paddle.log(pred_xt_1_probs + eps))) + .reshape([onehot_xt.shape[0], -1]).sum(axis=-1) + ) + return self.default_loss_scale * kl_loss diff --git a/structure_generation/configs/miad/README.md b/structure_generation/configs/miad/README.md new file mode 100644 index 00000000..a0360880 --- /dev/null +++ b/structure_generation/configs/miad/README.md @@ -0,0 +1,131 @@ +# MiAD + +[Mirage Atom Diffusion for De Novo Crystal Generation](https://arxiv.org/abs/2511.14426) + +## Abstract + +MiAD introduces Mirage Infusion, a mechanism that allows diffusion models to dynamically adjust the number of atoms in a crystal structure during the generation trajectory. By treating a variable number of atoms as "mirage" atoms (sentinel states), MiAD achieves state-of-the-art performance in generating stable, unique, and novel (S.U.N.) materials. It uses DiffCsp as the backbone denoising architecture. + +--- + +## Model Description + +### Overview + +A crystal is represented by three components within its unit cell: +- atom types: $A = (a_1,\ldots,a_N)$ +- fractional coordinates: $F = (f_1,\ldots,f_N),\; f_i \in [0,1)^3$ +- lattice matrix: $L \in \mathbb{R}^{3 \times 3}$ + +MiAD defines separate forward corruption processes for $(L, F, A)$ and trains a CSPNet-based denoiser to reverse them. During sampling, mirage atoms (type 0) can be transformed into real elements or discarded, allowing the model to adjust atom counts dynamically. + +### Method + +#### 1) Lattice diffusion + +The lattice is diffused with a standard DDPM Gaussian process. Supports Flow Matching as an alternative: + +$$ +L_t = \sqrt{\bar{\alpha}_t}\,L_0 + \sqrt{1 - \bar{\alpha}_t}\,\epsilon,\quad \epsilon \sim \mathcal{N}(0, I) +$$ + +#### 2) Fractional-coordinate diffusion on a torus + +Fractional coordinates live on a 3D torus $[0,1)^3$. Wrapped Normal noise is used: + +$$ +x_t = (x_0 + \sigma(t)\,\epsilon) \bmod 1,\quad \epsilon \sim \mathcal{N}(0, I) +$$ + +Also supports Periodic Flow Matching. + +#### 3) Atom-type diffusion + +Atom types use D3PM (uniform transition + cosine schedule) or DDPM with one-hot encoding. + +#### 4) Mirage Infusion + +During sampling, atoms with type 0 are treated as mirage atoms. The model dynamically decides which mirage atoms should materialize into real elements, allowing the final structure to have fewer atoms than the initial maximum. This is the core innovation enabling variable-atom-count generation. + +--- + +## Dataset Description + +MiAD is trained and evaluated on the MP-20 benchmark. + +#### MP-20 split + +| Dataset | Train | Val | Test | +| --- | --- | --- | --- | +| [MP-20](https://paddle-org.bj.bcebos.com/paddlematerial/datasets/mp_20/mp_20.zip) | 27136 | 9047 | 9046 | + +Extract to `./data/mp_20/` so that CSV files are at `./data/mp_20/train.csv`, `./data/mp_20/val.csv`, and `./data/mp_20/test.csv`. + +--- + +## Results + +| Model | Dataset | Config | Checkpoint / Log | +| --- | --- | --- | --- | +| miad_mp20 | mp20 | [miad_mp20.yaml](miad_mp20.yaml) | [checkpoint / log](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/MiAD/miad_mp20.zip) | + +--- + +## Command + +### Training + +```bash +# single-gpu training +python structure_generation/train.py -c structure_generation/configs/miad/miad_mp20.yaml +``` + +### Validation + +```bash +# Adjust program behavior on the fly using command-line parameters without modifying the configuration file directly. +# Example: --Global.do_eval=True +python structure_generation/train.py -c structure_generation/configs/miad/miad_mp20.yaml Global.do_eval=True Global.do_train=False Global.do_test=False Trainer.pretrained_model_path='path/to/model.pdparams' +``` + +### Testing + +```bash +# Evaluate the model on the test dataset. +python structure_generation/train.py -c structure_generation/configs/miad/miad_mp20.yaml Global.do_eval=False Global.do_train=False Global.do_test=True Trainer.pretrained_model_path='path/to/model.pdparams' +``` + +### Sample + +```bash +# This command is used to sample crystal structures using a trained model. +# Mode 1: Use a pre-trained model (downloads automatically). +# Mode 2: Use a custom configuration file and checkpoint. +# Results are saved to the folder specified by --save_path (default: results). + +# Mode 1: pre-trained model, sample by number of atoms +python structure_generation/sample.py --model_name='miad_mp20' --weights_name='miad_mp20.pdparams' --save_path='result_miad/' --mode='by_num_atoms' --num_atoms=20 + +# Mode 1: pre-trained model, sample by dataloader (reads test.csv) +python structure_generation/sample.py --model_name='miad_mp20' --weights_name='miad_mp20.pdparams' --save_path='result_miad/' --mode='by_dataloader' + +# Mode 2: custom config + checkpoint, sample by number of atoms +python structure_generation/sample.py --config_path='structure_generation/configs/miad/miad_mp20.yaml' --checkpoint_path='./output/miad_mp20/checkpoints/latest.pdparams' --save_path='result_miad/' --mode='by_num_atoms' --num_atoms=20 + +# Mode 2: custom config + checkpoint, sample by dataloader +python structure_generation/sample.py --config_path='structure_generation/configs/miad/miad_mp20.yaml' --checkpoint_path='./output/miad_mp20/checkpoints/latest.pdparams' --save_path='result_miad/' --mode='by_dataloader' +``` + +--- + +## Citation + +```bibtex +@article{okhotin2025miad, + title={MiAD: Mirage Atom Diffusion for De Novo Crystal Generation}, + author={Andrey Okhotin, Maksim Nakhodnov, Nikita Kazeev, Andrey E Ustyuzhanin, Dmitry Vetrov}, + journal={arXiv preprint arXiv:2511.14426}, + year={2025}, + url={https://arxiv.org/abs/2511.14426} +} +``` diff --git a/structure_generation/configs/miad/miad_mp20.yaml b/structure_generation/configs/miad/miad_mp20.yaml new file mode 100644 index 00000000..8446a00f --- /dev/null +++ b/structure_generation/configs/miad/miad_mp20.yaml @@ -0,0 +1,163 @@ +Global: + do_train: True + do_eval: False + do_test: False + num_train_timesteps: 1000 + +Trainer: + max_epochs: 2 + seed: 42 + output_dir: ./output/miad_mp20 + save_freq: 100 + log_freq: 10 + start_eval_epoch: 1 + eval_freq: 1 + pretrained_model_path: null + resume_from_checkpoint: null + use_amp: False + amp_level: 'O1' + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "loss" + greater_is_better: False + compute_metric_during_train: False + metric_strategy_during_eval: 'step' + use_visualdl: False + use_wandb: False + use_tensorboard: False + +Model: + __class_name__: MiAD + __init_params__: + model_cfg: + hidden_dim: 512 + latent_dim: 256 + num_layers: 6 + max_atoms: 100 + act_fn: silu + dis_emb: sin + num_freqs: 128 + edge_style: fc + ln: true + ip: true + smooth: true + pred_type: true + diffusion_cfg: + method: DiffCSP + task: gen_mp20 + cont_time: false + num_steps: 1000 + lat_diffusion: + method: ddpm + scheduler_cfg: + __class_name__: DDPMScheduler + __init_params__: + num_train_timesteps: 1000 + beta_schedule: squaredcos_cap_v2 + frac_diffusion: + method: wrapped_normal + scheduler_cfg: + __class_name__: ScoreSdeVeSchedulerWrapped + __init_params__: + num_train_timesteps: 1000 + sigma_min: 0.005 + sigma_max: 0.5 + sampling_eps: 0.001 + type_diffusion: + method: d3pm + +Optimizer: + __class_name__: Adam + __init_params__: + beta1: 0.9 + beta2: 0.999 + lr: + __class_name__: ReduceOnPlateau + __init_params__: + learning_rate: 0.001 + factor: 0.6 + by_epoch: True + patience: 30 + min_lr: 0.0001 + indicator: "train_loss" + indicator_name: 'loss' + +Dataset: + train: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/train.csv" + build_structure_cfg: + format: cif_str + num_cpus: 10 + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 4 + val: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/val.csv" + build_structure_cfg: + format: cif_str + num_cpus: 10 + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 4 + loader: + collate_fn: DefaultCollator + test: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/test.csv" + build_structure_cfg: + format: cif_str + num_cpus: 10 + sampler: + __class_name__: DistributedBatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 4 + loader: + collate_fn: DefaultCollator + +Sample: + data: + dataset: + __class_name__: MP20Dataset + __init_params__: + path: "./data/mp_20/test.csv" + build_structure_cfg: + format: cif_str + num_cpus: 10 + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 4 + loader: + collate_fn: DefaultCollator + build_structure_cfg: + format: array + niggli: False + metrics: + __class_name__: CSPMetric + __init_params__: + gt_file_path: "./data/mp_20/test.csv" + model_sample_params: + num_inference_steps: 1000 diff --git a/test/miad/__init__.py b/test/miad/__init__.py new file mode 100644 index 00000000..290f972c --- /dev/null +++ b/test/miad/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/test/miad/test_miad.py b/test/miad/test_miad.py new file mode 100644 index 00000000..2312e29c --- /dev/null +++ b/test/miad/test_miad.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import paddle +from omegaconf import OmegaConf + +from ppmat.datasets.collate_fn import DefaultCollator +from ppmat.datasets.mp20_dataset import MP20Dataset +from ppmat.models import build_model +from ppmat.models.miad.miad import _extract_x0 +from ppmat.models.miad.miad import MiAD + +TINY_MODEL_CFG = { + "hidden_dim": 64, + "latent_dim": 32, + "num_layers": 2, + "max_atoms": 100, + "act_fn": "silu", + "dis_emb": "sin", + "num_freqs": 10, + "edge_style": "fc", + "ln": False, + "ip": True, + "smooth": True, + "pred_type": True, +} + +TINY_DIFFUSION_CFG = { + "method": "DiffCSP", + "task": "gen_mp20", + "cont_time": False, + "num_steps": 10, + "time_embed_dim": 32, + "lat_diffusion": { + "method": "ddpm", + "scheduler_cfg": { + "__class_name__": "DDPMScheduler", + "__init_params__": { + "num_train_timesteps": 10, + "beta_schedule": "squaredcos_cap_v2", + }, + }, + }, + "frac_diffusion": { + "method": "wrapped_normal", + "scheduler_cfg": { + "__class_name__": "ScoreSdeVeSchedulerWrapped", + "__init_params__": { + "num_train_timesteps": 10, + "sigma_min": 0.005, + "sigma_max": 0.5, + "sampling_eps": 0.001, + }, + }, + }, + "type_diffusion": {"method": "d3pm"}, +} + + +def _make_fake_batch(batch_size=2, atoms_per_crystal=5): + num_atoms = paddle.full([batch_size], atoms_per_crystal, dtype="int64") + total_atoms = batch_size * atoms_per_crystal + batch_idx = paddle.concat( + [paddle.full([atoms_per_crystal], i, dtype="int64") for i in range(batch_size)] + ) + lattices = paddle.randn([batch_size, 3, 3], dtype="float32") + frac_coords = paddle.rand([total_atoms, 3], dtype="float32") + atom_types = paddle.randint(1, 10, [total_atoms], dtype="int64") + return { + "x0": [lattices, frac_coords, atom_types], + "batch_size": batch_size, + "num_atoms": num_atoms, + "batch_idx": batch_idx, + "atom_types": atom_types, + } + + +class MiADSmokeTest(unittest.TestCase): + """Self-contained smoke tests for MiAD (no weights, no external files).""" + + @classmethod + def setUpClass(cls): + paddle.seed(42) + + def test_forward_smoke(self): + model = MiAD(model_cfg=TINY_MODEL_CFG, diffusion_cfg=TINY_DIFFUSION_CFG) + model.eval() + batch = _make_fake_batch() + with paddle.no_grad(): + output = model(batch) + self.assertIn("loss_dict", output) + self.assertIn("loss", output["loss_dict"]) + loss = output["loss_dict"]["loss"] + self.assertTrue(paddle.isfinite(loss)) + + def test_sample_output_format(self): + model = MiAD(model_cfg=TINY_MODEL_CFG, diffusion_cfg=TINY_DIFFUSION_CFG) + model.eval() + batch_data = {"num_atoms": paddle.to_tensor([5, 7], dtype="int64")} + result = model.sample(batch_data, num_inference_steps=5) + self.assertIn("result", result) + self.assertEqual(len(result["result"]), 2) + for entry in result["result"]: + for key in ("num_atoms", "atom_types", "frac_coords", "lattice"): + self.assertIn(key, entry) + self.assertEqual(entry["frac_coords"].shape[-1], 3) + + +class MiADConfigTest(unittest.TestCase): + """Test MiAD construction via build_model from yaml config.""" + + def test_config_load(self): + config = OmegaConf.load("structure_generation/configs/miad/miad_mp20.yaml") + config = OmegaConf.to_container(config, resolve=True) + self.assertIn("Model", config) + self.assertEqual(config["Model"]["__class_name__"], "MiAD") + self.assertIn("diffusion_cfg", config["Model"]["__init_params__"]) + self.assertIn("model_cfg", config["Model"]["__init_params__"]) + + def test_build_model_path(self): + config = OmegaConf.load("structure_generation/configs/miad/miad_mp20.yaml") + config = OmegaConf.to_container(config, resolve=True) + model = build_model(config["Model"]) + self.assertIsInstance(model, MiAD) + self.assertIsInstance(model, paddle.nn.Layer) + + def test_train_path(self): + config = OmegaConf.load("structure_generation/configs/miad/miad_mp20.yaml") + config = OmegaConf.to_container(config, resolve=True) + model = build_model(config["Model"]) + model.eval() + batch = _make_fake_batch() + with paddle.no_grad(): + output = model(batch) + self.assertIn("loss_dict", output) + self.assertIn("loss", output["loss_dict"]) + self.assertTrue(paddle.isfinite(output["loss_dict"]["loss"])) + + def test_sample_path(self): + config = OmegaConf.load("structure_generation/configs/miad/miad_mp20.yaml") + config = OmegaConf.to_container(config, resolve=True) + model = build_model(config["Model"]) + model.eval() + batch_data = {"num_atoms": paddle.to_tensor([5, 7], dtype="int64")} + result = model.sample(batch_data, num_inference_steps=5) + self.assertIn("result", result) + self.assertEqual(len(result["result"]), 2) + for entry in result["result"]: + self.assertIn("num_atoms", entry) + self.assertIn("lattice", entry) + + +class MiADDatasetTest(unittest.TestCase): + """Dataset smoke test for MiAD.""" + + @classmethod + def setUpClass(cls): + cls.dataset = MP20Dataset( + path="./data/mp_20/test.csv", + build_structure_cfg={"format": "cif_str", "num_cpus": 1}, + ) + + def test_dataset_load(self): + self.assertGreater(len(self.dataset), 0) + + def test_dataset_sample_fields(self): + sample = self.dataset[0] + self.assertIn("structure_array", sample) + sa = sample["structure_array"] + self.assertIn("frac_coords", sa) + self.assertIn("atom_types", sa) + self.assertIn("lattice", sa) + self.assertIn("num_atoms", sa) + + def test_collate_fn(self): + collator = DefaultCollator() + samples = [self.dataset[i] for i in range(min(4, len(self.dataset)))] + batch = collator(samples) + batch = _extract_x0(batch) + self.assertIn("x0", batch) + self.assertIn("batch_size", batch) + self.assertIn("num_atoms", batch) + self.assertIn("batch_idx", batch) + x0 = batch["x0"] + self.assertEqual(len(x0), 3) + self.assertEqual(x0[0].ndim, 3) + self.assertEqual(x0[1].ndim, 2) + self.assertEqual(x0[2].ndim, 1) + self.assertEqual(x0[0].shape[-1], 3) + self.assertEqual(x0[1].shape[-1], 3) + + +if __name__ == "__main__": + unittest.main()