From 0a4e755e8966dcb7998b49edde20c6e6ef6067de Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Wed, 27 May 2026 16:30:57 +0800 Subject: [PATCH 01/10] add maid model --- ppmat/datasets/collate_fn.py | 1 + ppmat/metrics/__init__.py | 2 + ppmat/metrics/prerelax_chgnet.py | 248 ++++++++++ ppmat/metrics/sun_metric.py | 423 ++++++++++++++++++ ppmat/models/__init__.py | 3 + ppmat/models/miad/__init__.py | 58 +++ ppmat/models/miad/collate.py | 237 ++++++++++ ppmat/models/miad/crystal_diffusion.py | 356 +++++++++++++++ ppmat/models/miad/default_configs.py | 88 ++++ ppmat/models/miad/diffusion_utils.py | 96 ++++ ppmat/models/miad/frac_diffusion.py | 183 ++++++++ ppmat/models/miad/graph_utils.py | 52 +++ ppmat/models/miad/lattice_diffusion.py | 253 +++++++++++ ppmat/models/miad/miad.py | 256 +++++++++++ ppmat/models/miad/miad_cspnet.py | 133 ++++++ ppmat/models/miad/type_diffusion.py | 222 +++++++++ ppmat/schedulers/__init__.py | 6 + ppmat/schedulers/miad_schedulers.py | 146 ++++++ structure_generation/configs/miad/README.md | 77 ++++ .../configs/miad/miad_mp20.yaml | 154 +++++++ test/test_miad.py | 114 +++++ 21 files changed, 3108 insertions(+) create mode 100644 ppmat/metrics/prerelax_chgnet.py create mode 100644 ppmat/metrics/sun_metric.py create mode 100644 ppmat/models/miad/__init__.py create mode 100644 ppmat/models/miad/collate.py create mode 100644 ppmat/models/miad/crystal_diffusion.py create mode 100644 ppmat/models/miad/default_configs.py create mode 100644 ppmat/models/miad/diffusion_utils.py create mode 100644 ppmat/models/miad/frac_diffusion.py create mode 100644 ppmat/models/miad/graph_utils.py create mode 100644 ppmat/models/miad/lattice_diffusion.py create mode 100644 ppmat/models/miad/miad.py create mode 100644 ppmat/models/miad/miad_cspnet.py create mode 100644 ppmat/models/miad/type_diffusion.py create mode 100644 ppmat/schedulers/miad_schedulers.py create mode 100644 structure_generation/configs/miad/README.md create mode 100644 structure_generation/configs/miad/miad_mp20.yaml create mode 100644 test/test_miad.py diff --git a/ppmat/datasets/collate_fn.py b/ppmat/datasets/collate_fn.py index 9073af4c..69b68e19 100644 --- a/ppmat/datasets/collate_fn.py +++ b/ppmat/datasets/collate_fn.py @@ -30,6 +30,7 @@ from ppmat.datasets.custom_data_type import ConcatNumpyWarper from ppmat.datasets.geometric_data_type.batch import Batch from ppmat.datasets.geometric_data_type.data import Data +from ppmat.models.miad.collate import MiADCollator # noqa: F401 class DefaultCollator(object): diff --git a/ppmat/metrics/__init__.py b/ppmat/metrics/__init__.py index a0e3fb75..088b1975 100644 --- a/ppmat/metrics/__init__.py +++ b/ppmat/metrics/__init__.py @@ -18,6 +18,7 @@ from ppmat.metrics.csp_metric import CSPMetric from ppmat.metrics.diffnmr_streaming_adapter import DiffNMRStreamingAdapter +from ppmat.metrics.sun_metric import SUNMetric __all__ = [ "build_metric", @@ -25,6 +26,7 @@ "DiffNMRStreamingAdapter", # "DiffNMRMetric", # "NLL", "CrossEntropyMetric", "SumExceptBatchMetric", "SumExceptBatchKL", + "SUNMetric", ] diff --git a/ppmat/metrics/prerelax_chgnet.py b/ppmat/metrics/prerelax_chgnet.py new file mode 100644 index 00000000..2b090833 --- /dev/null +++ b/ppmat/metrics/prerelax_chgnet.py @@ -0,0 +1,248 @@ +# Copyright (c) 2025 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. + +""" +Pre-relaxation module using ppmat's built-in CHGNet for energy prediction. + +Provides structure relaxation and energy calculation needed for the +Stability metric in S.U.N. evaluation. Uses ppmat's CHGNet model +(Paddle implementation) -- no external PyTorch dependencies required. +""" + +import logging +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple + +import numpy as np +from pymatgen.core import Structure +from tqdm import tqdm + +logger = logging.getLogger(__name__) + +# Lazy-loaded globals to avoid import cost at module level +_chgnet_model = None +_graph_converter = None + + +def _load_chgnet(model_name: str = "chgnet_mptrj"): + """Lazy-load CHGNet model and graph converter.""" + 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", + batch_size: int = 1, +) -> List[Optional[float]]: + """Predict energy per atom for each structure using CHGNet. + + Args: + structures: List of pymatgen Structures (None for invalid). + model_name: CHGNet model name in ppmat registry. + batch_size: Prediction batch size. + + Returns: + List of energy per atom (eV/atom), None for invalid structures. + """ + 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 relax_structure( + structure: Structure, + model_name: str = "chgnet_mptrj", + steps: int = 200, + fmax: float = 0.05, + step_size: float = 0.02, +) -> Tuple[Optional[Structure], Optional[float], int]: + """Relax a crystal structure using gradient-based optimization with CHGNet. + + Performs iterative structure relaxation by computing forces from CHGNet + energy predictions and updating atomic positions accordingly. + + Args: + structure: pymatgen Structure to relax. + model_name: CHGNet model name in ppmat registry. + steps: Maximum number of relaxation steps. + fmax: Force convergence threshold (eV/A). + step_size: Position update step size. + + Returns: + Tuple of (relaxed_structure, energy_per_atom, n_steps). + None values if relaxation fails. + """ + model, converter = _load_chgnet(model_name) + + try: + import paddle + + current_struct = structure.copy() + converged_step = steps + + for step in range(steps): + graph = converter(current_struct) + pred = model.predict(graph) + e_per_atom = float(np.array(pred["energy_per_atom"]).flatten()[0]) + + forces = np.array(pred["force"]) + max_force = np.max(np.abs(forces)) + + if max_force < fmax: + converged_step = step + 1 + break + + # Simple gradient descent on positions + frac_coords = current_struct.frac_coords + lattice_matrix = current_struct.lattice.matrix + + # Force in Cartesian -> displacement in fractional + cart_displacements = -step_size * forces / (max_force + 1e-8) + frac_displacements = cart_displacements @ np.linalg.inv(lattice_matrix) + + new_frac = frac_coords + frac_displacements + # Wrap to [0, 1) + new_frac = new_frac % 1.0 + + current_struct = Structure( + current_struct.lattice, + current_struct.species, + new_frac, + coords_are_cartesian=False, + ) + + # Final energy + graph = converter(current_struct) + pred = model.predict(graph) + final_energy = float(np.array(pred["energy_per_atom"]).flatten()[0]) + + return current_struct, final_energy, converged_step + + except Exception as e: + logger.debug(f"Relaxation failed: {e}") + return None, None, 0 + + +def predict_energies_with_relaxation( + structures: List[Optional[Structure]], + model_name: str = "chgnet_mptrj", + steps: int = 200, + fmax: float = 0.05, + relax: bool = True, +) -> Dict[str, List]: + """Predict energies for structures, optionally with pre-relaxation. + + This is the main entry point for S.U.N. Stability computation. + + Args: + structures: List of pymatgen Structures (None for invalid). + model_name: CHGNet model name in ppmat registry. + steps: Maximum relaxation steps. + fmax: Force convergence threshold (eV/A). + relax: Whether to relax structures before energy evaluation. + + Returns: + Dict with: + - 'energy_gen': energy per atom before relaxation + - 'energy_relaxed': energy per atom after relaxation (or same as gen) + - 'relaxed_structures': relaxed Structures (or originals) + - 'num_sites': number of atoms per structure + """ + energy_gen = [] + energy_relaxed = [] + relaxed_structures = [] + num_sites = [] + + for struct in tqdm( + structures, + desc=f"{'Relaxing' if relax else 'Evaluating'} structures (CHGNet)", + ): + if struct is None: + energy_gen.append(None) + energy_relaxed.append(None) + relaxed_structures.append(None) + num_sites.append(None) + continue + + try: + n_sites = struct.num_sites + model, converter = _load_chgnet(model_name) + + # Initial energy + graph = converter(struct) + pred = model.predict(graph) + e_gen = float(np.array(pred["energy_per_atom"]).flatten()[0]) * n_sites + + if relax: + relaxed, e_relax, n_steps = relax_structure( + struct, model_name=model_name, steps=steps, fmax=fmax + ) + if relaxed is not None: + energy_gen.append(e_gen / n_sites) + energy_relaxed.append(e_relax) + relaxed_structures.append(relaxed) + num_sites.append(n_sites) + else: + energy_gen.append(e_gen / n_sites) + energy_relaxed.append(e_gen / n_sites) + relaxed_structures.append(struct) + num_sites.append(n_sites) + else: + energy_gen.append(e_gen / n_sites) + energy_relaxed.append(e_gen / n_sites) + relaxed_structures.append(struct) + num_sites.append(n_sites) + + except Exception as e: + logger.debug(f"Structure evaluation failed: {e}") + energy_gen.append(None) + energy_relaxed.append(None) + relaxed_structures.append(None) + num_sites.append(None) + + return { + "energy_gen": energy_gen, + "energy_relaxed": energy_relaxed, + "relaxed_structures": relaxed_structures, + "num_sites": num_sites, + } diff --git a/ppmat/metrics/sun_metric.py b/ppmat/metrics/sun_metric.py new file mode 100644 index 00000000..8e968814 --- /dev/null +++ b/ppmat/metrics/sun_metric.py @@ -0,0 +1,423 @@ +# Copyright (c) 2025 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. + +""" +S.U.N. (Stability, Uniqueness, Novelty) metrics for crystal structure generation. + +Implements the evaluation pipeline from MiAD paper for assessing generated +crystal structures across three dimensions: +- Stability: energy above hull below a threshold +- Uniqueness: no duplicates among generated structures +- Novelty: not present in the training set + +Reference: https://arxiv.org/abs/2511.14426 +""" + +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 + +logger = logging.getLogger(__name__) + + +def _composition_hash(structure: Structure) -> str: + """Hash a structure by its sorted atomic numbers for quick grouping.""" + return str(sorted(list(structure.atomic_numbers))) + + +def _structures_from_cif_strings(cif_strings: List[str]) -> List[Optional[Structure]]: + """Parse CIF strings into pymatgen Structures.""" + structures = [] + for cif in cif_strings: + try: + structures.append(Structure.from_str(cif, fmt="cif")) + except Exception: + structures.append(None) + return structures + + +def _structures_from_dicts(structure_dicts: List[dict]) -> List[Optional[Structure]]: + """Build Structures from dict lists with lattice, atom_types, frac_coords.""" + from pymatgen.core import Lattice + + structures = [] + for d in structure_dicts: + try: + lattice = Lattice(d["lattice"]) + struct = Structure( + lattice, + d["atom_types"], + d["frac_coords"], + coords_are_cartesian=False, + ) + structures.append(struct) + except Exception: + structures.append(None) + return structures + + +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]: + """Compute uniqueness of generated structures. + + A structure is unique if it does not match any previously seen structure + with the same composition. + + Args: + structures: List of pymatgen Structures (None for invalid). + stol: Site tolerance for StructureMatcher. + angle_tol: Angle tolerance for StructureMatcher. + ltol: Length tolerance for StructureMatcher. + attempt_supercell: Whether to attempt supercell matching. + symmetric: Whether to use symmetric matching. + + Returns: + List of booleans indicating uniqueness. + """ + matcher = StructureMatcher( + stol=stol, angle_tol=angle_tol, ltol=ltol, + attempt_supercell=attempt_supercell, + ) + seen = defaultdict(list) + uniqueness = [] + + for struct in tqdm(structures, desc="Computing uniqueness"): + if struct is None: + uniqueness.append(False) + continue + h = _composition_hash(struct) + if h not in seen: + uniqueness.append(True) + else: + for existing in seen[h]: + if matcher.fit(struct, existing, symmetric=symmetric): + uniqueness.append(False) + break + else: + uniqueness.append(True) + seen[h].append(struct) + + return uniqueness + + +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]: + """Compute novelty of generated structures against a reference set. + + A structure is novel if it does not match any structure in the reference + set (typically the training set). + + Args: + structures: List of pymatgen Structures (None for invalid). + reference_structures: Reference structures (training set). + stol: Site tolerance for StructureMatcher. + angle_tol: Angle tolerance for StructureMatcher. + ltol: Length tolerance for StructureMatcher. + attempt_supercell: Whether to attempt supercell matching. + symmetric: Whether to use symmetric matching. + + Returns: + List of booleans indicating novelty. + """ + matcher = StructureMatcher( + stol=stol, angle_tol=angle_tol, ltol=ltol, + attempt_supercell=attempt_supercell, + ) + + reference_by_comp = defaultdict(list) + for ref in reference_structures: + h = _composition_hash(ref) + reference_by_comp[h].append(ref) + + novelty = [] + for struct in tqdm(structures, desc="Computing novelty"): + if struct is None: + novelty.append(False) + continue + h = _composition_hash(struct) + if h not in reference_by_comp: + novelty.append(True) + else: + for ref in reference_by_comp[h]: + if matcher.fit(struct, ref, symmetric=symmetric): + novelty.append(False) + break + else: + novelty.append(True) + + return novelty + + +def compute_stability( + energy_above_hull: List[Optional[float]], + threshold: float = 0.0, +) -> List[bool]: + """Compute stability from energy above hull values. + + Args: + energy_above_hull: Energy above hull per atom (None for invalid). + threshold: Energy threshold in eV/atom. 0.0 for stable, 0.1 for metastable. + + Returns: + List of booleans indicating stability. + """ + 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]: + """Compute combined S.U.N. metrics. + + S.U.N. = structure is Stable AND Unique AND Novel AND non-trivial + (more than one element type). + + Args: + stability: Stability flags. + uniqueness: Uniqueness flags. + novelty: Novelty flags. + structures: Original structures (for composition check). + min_elements: Minimum number of element types for non-trivial. + + Returns: + Dictionary with metric rates. + """ + n = len(stability) + assert len(uniqueness) == n and len(novelty) == n + + 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: + """S.U.N. metric for evaluating generated crystal structures. + + Computes Stability (energy above hull), Uniqueness (no duplicates), + Novelty (not in training set), and their combination. + + Args: + gt_file_path: Path to ground truth CSV (with 'cif' column) for novelty ref. + stol: Site tolerance for structure matching. + angle_tol: Angle tolerance for structure matching. + ltol: Length tolerance for structure matching. + stability_threshold: Energy above hull threshold in eV/atom. + attempt_supercell: Whether to attempt supercell matching. + """ + + 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, + chgnet_relax: bool = True, + chgnet_relax_steps: int = 200, + ): + 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.chgnet_relax = chgnet_relax + self.chgnet_relax_steps = chgnet_relax_steps + self._reference_structures = None + + def _load_reference(self): + """Load reference structures from ground truth CSV.""" + 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( + Structure.from_str(cif_str, fmt="cif") + ) + 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]: + """Run the full S.U.N. evaluation pipeline. + + Args: + generated: Generated structures. Can be: + - List of dicts with 'lattice', 'atom_types', 'frac_coords' + - List of CIF strings + - DataFrame with a 'cif' column + energy_above_hull: Optional pre-computed energy above hull values. + If None, stability is not computed (all marked False). + + Returns: + Dictionary with S.U.N. metrics. + """ + # Parse structures + if isinstance(generated, pd.DataFrame): + if "cif" in generated.columns: + structures = _structures_from_cif_strings(generated["cif"].tolist()) + elif "structure" in generated.columns: + structures = generated["structure"].tolist() + else: + raise ValueError("DataFrame must have 'cif' or 'structure' column") + elif isinstance(generated, list) and len(generated) > 0: + if isinstance(generated[0], str): + structures = _structures_from_cif_strings(generated) + elif isinstance(generated[0], dict): + structures = _structures_from_dicts(generated) + else: + structures = generated + else: + raise ValueError("generated must be a list or DataFrame") + + n = len(structures) + logger.info(f"Evaluating {n} generated structures") + + # Uniqueness + uniqueness = compute_uniqueness( + structures, + stol=self.stol, angle_tol=self.angle_tol, ltol=self.ltol, + attempt_supercell=self.attempt_supercell, + ) + + # Novelty + novelty_list = [True] * n # default: all novel + 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, + ) + + # Stability + if energy_above_hull is not None: + stability = compute_stability( + energy_above_hull, threshold=self.stability_threshold + ) + elif self.use_chgnet: + from ppmat.metrics.prerelax_chgnet import predict_energy + + 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 provided and use_chgnet=False, " + "stability metrics will be 0%" + ) + + # Combined S.U.N. + results = compute_sun(stability, uniqueness, novelty_list, structures) + + if self.use_chgnet and energies is not None: + valid_energies = [e for e in energies if e is not None] + if valid_energies: + results["avg_energy_per_atom"] = float(np.mean(valid_energies)) + + # Also compute metastable variant + if energy_above_hull is not None: + metastability = compute_stability( + energy_above_hull, threshold=0.1 + ) + msun_results = compute_sun( + metastability, uniqueness, novelty_list, structures + ) + results["metastable_rate"] = msun_results["stability_rate"] + results["msun_rate"] = msun_results["sun_rate"] + results["msun_count"] = msun_results["sun_count"] + + # Log results + logger.info( + f"S.U.N. Results: " + f"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 95d73232..f6295669 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.utils import download from ppmat.utils import logger from ppmat.utils import save_load @@ -67,6 +68,7 @@ "DiffNMR", "InfGCN", "MatENO", + "MiAD", ] # Warning: The key of the dictionary must be consistent with the file name of the value @@ -111,6 +113,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", } diff --git a/ppmat/models/miad/__init__.py b/ppmat/models/miad/__init__.py new file mode 100644 index 00000000..6154349b --- /dev/null +++ b/ppmat/models/miad/__init__.py @@ -0,0 +1,58 @@ +# 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.collate import CrystalBatch +from ppmat.models.miad.collate import MiADCollator +from ppmat.models.miad.collate import create_miad_dataloader +from ppmat.models.miad.collate import create_sampling_batch +from ppmat.models.miad.crystal_diffusion import CrystalGen as MiadCrystalGen +from ppmat.models.miad.crystal_diffusion import DiffCSP as MiadDiffCSP +from ppmat.models.miad.crystal_diffusion import init_diffusion +from ppmat.models.miad.crystal_diffusion import parse_batch +from ppmat.models.miad.diffusion_utils import TimeDistribution as MiadTimeDistribution +from ppmat.models.miad.frac_diffusion import PFM as MiadPFM +from ppmat.models.miad.frac_diffusion import WrappedNormal as MiadWrappedNormal +from ppmat.models.miad.lattice_diffusion import DDPM as MiadDDPM +from ppmat.models.miad.lattice_diffusion import FM as MiadFM +from ppmat.models.miad.lattice_diffusion import FM_LenAng as MiadFM_LenAng +from ppmat.models.miad.miad import MiAD +from ppmat.models.miad.miad_cspnet import CSPNet as MiadCSPNet +from ppmat.models.miad.type_diffusion import D3PM as MiadD3PM +from ppmat.models.miad.type_diffusion import DDPM_onehot as MiadDDPM_onehot +from ppmat.schedulers.miad_schedulers import scheduler + +__all__ = [ + # Model + "MiAD", + "MiadCSPNet", + "MiadCrystalGen", + "MiadDiffCSP", + "init_diffusion", + "parse_batch", + # Diffusion components + "MiadDDPM", + "MiadFM", + "MiadFM_LenAng", + "MiadWrappedNormal", + "MiadPFM", + "MiadDDPM_onehot", + "MiadD3PM", + "MiadTimeDistribution", + "scheduler", + # Collate + "MiADCollator", + "CrystalBatch", + "create_miad_dataloader", + "create_sampling_batch", +] diff --git a/ppmat/models/miad/collate.py b/ppmat/models/miad/collate.py new file mode 100644 index 00000000..5b7ba9bc --- /dev/null +++ b/ppmat/models/miad/collate.py @@ -0,0 +1,237 @@ +# 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 typing import Any +from typing import Dict +from typing import List +from typing import Optional + +import numpy as np +import paddle + +from ppmat.datasets.custom_data_type import ConcatData + + +def _extract_concat_data(data): + if data is None: + return None + if isinstance(data, ConcatData): + return np.array(data.data) + return data + + +class CrystalBatch: + """Minimal batch object mimicking torch_geometric Batch interface.""" + + def __init__( + self, + num_atoms: paddle.Tensor, + atom_types: paddle.Tensor, + batch: paddle.Tensor, + frac_coords: Optional[paddle.Tensor] = None, + lattice: Optional[paddle.Tensor] = None, + ): + self.num_atoms = num_atoms + self.atom_types = atom_types + self.batch = batch + self.frac_coords = frac_coords + self.lattice = lattice + + def cuda(self, device: str = "gpu"): + place = paddle.CUDAPlace(0) if device == "gpu" else paddle.CPUPlace() + self.num_atoms = self.num_atoms.to(place) + self.atom_types = self.atom_types.to(place) + self.batch = self.batch.to(place) + if self.frac_coords is not None: + self.frac_coords = self.frac_coords.to(place) + if self.lattice is not None: + self.lattice = self.lattice.to(place) + return self + + +class MiADCollator: + """Collate function for MiAD model. + + Converts MP20Dataset batch to CrystalGen format. + """ + + def __init__( + self, + mirage_max_atoms: Optional[int] = None, + use_mirage: bool = False, + ): + self.mirage_max_atoms = mirage_max_atoms + self.use_mirage = use_mirage + self.N_m = mirage_max_atoms + + def __call__(self, batch_list: List[Dict]) -> Dict[str, Any]: + if not batch_list: + return {} + batch_size = len(batch_list) + all_frac_coords = [] + all_atom_types = [] + all_num_atoms = [] + all_lattices = [] + batch_idx_list = [] + + for i, sample_i in enumerate(batch_list): + if not isinstance(sample_i, dict): + continue + if "structure_array" in sample_i: + sa = sample_i["structure_array"] + elif "frac_coords" in sample_i: + sa = sample_i + else: + continue + + frac_coords = _extract_concat_data(sa.get("frac_coords")) + atom_types = _extract_concat_data(sa.get("atom_types")) + lattice = _extract_concat_data(sa.get("lattice")) + num_atoms = _extract_concat_data(sa.get("num_atoms")) + + if frac_coords is None or atom_types is None or lattice is None: + continue + + if frac_coords.ndim >= 2: + n_atoms = frac_coords.shape[0] + elif num_atoms is not None: + if isinstance(num_atoms, np.ndarray) and num_atoms.ndim > 0: + n_atoms = int(num_atoms.flatten()[0]) + else: + n_atoms = int(num_atoms) + else: + continue + + if frac_coords.ndim == 1: + frac_coords = frac_coords.reshape(-1, 3) + if atom_types.ndim > 1: + atom_types = atom_types.flatten() + + if lattice.ndim == 2: + lattice = lattice.reshape(1, 3, 3) + elif lattice.ndim == 1: + lattice = lattice.reshape(1, 3, 3) + + if self.N_m is not None and self.use_mirage and n_atoms < self.N_m: + frac_coords, atom_types = self._pad_mirage( + frac_coords, atom_types, n_atoms, self.N_m + ) + n_atoms = self.N_m + + all_frac_coords.append(frac_coords) + all_atom_types.append(atom_types) + all_lattices.append(lattice[0] if lattice.shape[0] == 1 else lattice) + all_num_atoms.append(n_atoms) + batch_idx_list.extend([i] * n_atoms) + + if not all_frac_coords: + return {} + + frac_coords_cat = np.concatenate(all_frac_coords, axis=0).astype("float32") + atom_types_cat = np.concatenate(all_atom_types, axis=0).astype("int64") + lattices_cat = np.stack(all_lattices, axis=0).astype("float32") + num_atoms_arr = np.array(all_num_atoms, dtype="int64") + batch_idx = np.array(batch_idx_list, dtype="int64") + + frac_coords_t = paddle.to_tensor(frac_coords_cat) + atom_types_t = paddle.to_tensor(atom_types_cat) + lattices_t = paddle.to_tensor(lattices_cat) + num_atoms_t = paddle.to_tensor(num_atoms_arr) + batch_idx_t = paddle.to_tensor(batch_idx) + + crystal_batch = CrystalBatch( + num_atoms=num_atoms_t, + atom_types=atom_types_t, + batch=batch_idx_t, + frac_coords=frac_coords_t, + lattice=lattices_t, + ) + + return { + "x0": [lattices_t, frac_coords_t, atom_types_t], + "batch_size": batch_size, + "num_atoms": num_atoms_t, + "batch_idx": batch_idx_t, + "atom_types": atom_types_t, + "batch": crystal_batch, + } + + @staticmethod + def _pad_mirage(frac_coords, atom_types, n_atoms, N_m): + pad_n = N_m - n_atoms + if pad_n > 0: + pad_fc = np.random.rand(pad_n, 3).astype("float32") + pad_at = np.zeros(pad_n, dtype="int64") + frac_coords = np.concatenate([frac_coords, pad_fc], axis=0) + atom_types = np.concatenate([atom_types, pad_at], axis=0) + return frac_coords, atom_types + + +def create_miad_dataloader( + dataset, + batch_size: int, + shuffle: bool = False, + num_workers: int = 0, + mirage_max_atoms: Optional[int] = None, + use_mirage: bool = False, + **kwargs, +) -> paddle.io.DataLoader: + collate_fn = MiADCollator( + mirage_max_atoms=mirage_max_atoms, + use_mirage=use_mirage, + ) + return paddle.io.DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + collate_fn=collate_fn, + **kwargs, + ) + + +def create_sampling_batch( + batch_size: int, + num_atoms: Optional[paddle.Tensor] = None, + mirage_max_atoms: Optional[int] = None, + device: str = "cpu", +) -> Dict[str, Any]: + if mirage_max_atoms is not None: + N_m = mirage_max_atoms + num_atoms = paddle.full([batch_size], N_m, dtype="int64") + total_atoms = batch_size * N_m + elif num_atoms is not None: + total_atoms = int(num_atoms.sum()) + else: + num_atoms = paddle.randint(5, 30, [batch_size], dtype="int64") + total_atoms = int(num_atoms.sum()) + + node2graph = paddle.repeat_interleave( + paddle.arange(batch_size, dtype="int64"), num_atoms + ) + + crystal_batch = CrystalBatch( + num_atoms=num_atoms, + atom_types=paddle.zeros([total_atoms], dtype="int64"), + batch=node2graph, + frac_coords=paddle.zeros([total_atoms, 3], dtype="float32"), + lattice=paddle.zeros([batch_size, 3, 3], dtype="float32"), + ) + + return { + "x0": [None, None, None], + "batch_size": batch_size, + "num_atoms": num_atoms, + "batch": crystal_batch, + } diff --git a/ppmat/models/miad/crystal_diffusion.py b/ppmat/models/miad/crystal_diffusion.py new file mode 100644 index 00000000..c29a2108 --- /dev/null +++ b/ppmat/models/miad/crystal_diffusion.py @@ -0,0 +1,356 @@ +# 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. + +""" +Crystal diffusion controllers for MiAD. +""" + +import os + +import numpy as np +import paddle + +from ppmat.models.common.time_embedding import SinusoidalTimeEmbeddings +from ppmat.models.miad.diffusion_utils import TimeDistribution +from ppmat.models.miad.diffusion_utils import mean_interleave +from ppmat.models.miad.frac_diffusion import PFM +from ppmat.models.miad.frac_diffusion import WrappedNormal +from ppmat.models.miad.lattice_diffusion import DDPM +from ppmat.models.miad.lattice_diffusion import FM +from ppmat.models.miad.lattice_diffusion import FM_LenAng +from ppmat.models.miad.type_diffusion import D3PM +from ppmat.models.miad.type_diffusion import DDPM_onehot + + +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 parse_batch(batch): + if "batch" in batch: + return { + "num_atoms": batch["batch"].num_atoms, + "batch_idx": batch["batch"].batch, + "atom_types": batch["batch"].atom_types, + "batch_size": batch["batch_size"], + } + if "batch_idx" in batch: + return { + "num_atoms": batch["num_atoms"], + "batch_idx": batch["batch_idx"], + "atom_types": batch["atom_types"], + "batch_size": batch.get("batch_size", len(batch["num_atoms"])), + } + raise ValueError("Cannot determine batch format: missing 'batch' or 'batch_idx'") + + +def init_diffusion(diffusion_config, logger): + if hasattr(diffusion_config, "default_config") and diffusion_config.default_config: + diffusion_config = _apply_default_config(diffusion_config.default_config) + + switch = { + "Default": CrystalGen, + "DiffCSP": DiffCSP, + } + method = diffusion_config.method + if method in switch: + return switch[method](diffusion_config, logger) + raise NotImplementedError(f"Diffusion method '{method}' not implemented") + + +def _apply_default_config(config_name): + from ppmat.models.miad.default_configs import DEFAULT_DIFFUSION_CONFIGS + + if config_name in DEFAULT_DIFFUSION_CONFIGS: + return DEFAULT_DIFFUSION_CONFIGS[config_name] + raise KeyError(f"Unknown default config: '{config_name}'") + + +class CrystalGen: + """Standard crystal generation with single-step reverse sampling.""" + + 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.time_distribution = TimeDistribution(self.num_steps, self.cont_time) + time_embed_dim = getattr(self.config, "time_embed_dim", 256) + self.time_embedding = SinusoidalTimeEmbeddings(time_embed_dim) + + # Lattice diffusion + switch_lat = { + "ddpm": DDPM, + "fm": FM, + "fm_lenang": FM_LenAng, + } + self.lat_diffusion = switch_lat[self.config.lat_diffusion.method](self.config) + + # Fractional coordinate diffusion + switch_frac = { + "wrapped_normal": WrappedNormal, + "pfm": PFM, + } + self.frac_diffusion = switch_frac[self.config.frac_diffusion.method]( + self.config + ) + + # Atom type diffusion (only for generation tasks) + self.gen_type = "gen" in self.config.task + if self.gen_type: + switch_type = { + "ddpm_onehot": DDPM_onehot, + "d3pm": D3PM, + } + self.type_diffusion = switch_type[self.config.type_diffusion.method]( + self.config + ) + else: + self.type_diffusion = None + + def forward_step_sample(self, x0, t, batch): + l0, f0, a0 = x0 + lt = self.lat_diffusion.forward_step_sample(l0, t[0], batch) + ft = self.frac_diffusion.forward_step_sample(f0, t[1], batch) + at = ( + self.type_diffusion.forward_step_sample(a0, t[1], batch) + if self.gen_type + else a0 + ) + return [lt, ft, at] + + def reverse_step_sample(self, xt, t, model, batch): + lt, ft, at = xt + batch["prediction"] = self.model_prediction(xt, t, model, batch) + l_pred, f_pred, a_pred = batch["prediction"] + lt_1 = self.lat_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) + ft_1 = self.frac_diffusion.reverse_step_sample(f_pred, ft, t[1], batch) + 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 _get_batch_info(self, batch): + try: + return parse_batch(batch) + except ValueError: + return self._normalize_batch(batch) + + def _normalize_batch(self, batch): + if "structure_array" in batch: + num_atoms_data = batch["structure_array"].get("num_atoms", None) + else: + num_atoms_data = batch.get("num_atoms", None) + + parsed = parse_num_atoms_to_per_crystal(num_atoms_data) + if parsed is None: + raise ValueError("Cannot determine num_atoms from batch") + 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") + + return { + "num_atoms": num_atoms, + "batch_idx": batch_idx, + "atom_types": atom_types, + "batch_size": batch_size, + } + + def prior_sample(self, batch): + batch_info = self._get_batch_info(batch) + return [ + self.lat_diffusion.prior_sample(batch), + self.frac_diffusion.prior_sample(batch), + ( + self.type_diffusion.prior_sample(batch) + if self.gen_type + else batch_info["atom_types"] + ), + ] + + def model_prediction(self, xt, t, model, batch): + batch_info = self._get_batch_info(batch) + lt, ft, at = xt + nn_pred = model( + t[0], + self.time_embedding(1000 * (t[0] / self.num_steps) + 1), + at, + ft, + lt, + batch_info["num_atoms"], + batch_info["batch_idx"], + ) + l_pred = nn_pred[0] + f_pred = nn_pred[1] + a_pred = nn_pred[2] if self.gen_type else None + return [l_pred, f_pred, a_pred] + + def get_x0_prediction(self, pred, xt, t, batch): + l_pred, f_pred, a_pred = pred + lt, ft, at = xt + return [ + self.lat_diffusion.get_x0_prediction(l_pred, lt, t[0], batch), + self.frac_diffusion.get_x0_prediction(f_pred, ft, t[1], batch), + ( + self.type_diffusion.get_x0_prediction( + a_pred, at, t[1], batch, x0_format="disc" + ) + if self.gen_type + else at.clone().detach() + ), + ] + + def train_step(self, batch, model, mode): + modifications = os.environ.get("MODIFICATIONS_FIELD", "") + + batch["t"] = self.time_distribution.sample(batch, mode) + batch["xt"] = self.forward_step_sample(batch["x0"], batch["t"], batch) + batch["prediction"] = self.model_prediction( + batch["xt"], batch["t"], model, batch + ) + + loss_lat = self.lat_diffusion.loss(batch) + loss_frac = self.frac_diffusion.loss(batch) + + if "insert_latloss_" in modifications: + lat_coef = float( + modifications.split("insert_latloss_")[1].split("_coef")[0] + ) + loss_lat = loss_lat * lat_coef + if "insert_fracloss_" in modifications: + frac_coef = float( + modifications.split("insert_fracloss_")[1].split("_coef")[0] + ) + loss_frac = loss_frac * frac_coef + + loss_lat_val = loss_lat.mean() + loss_frac_val = loss_frac.mean() + batch["loss"] = loss_lat_val + loss_frac_val + + if self.gen_type: + loss_type = self.type_diffusion.loss(batch) + if "insert_typeloss_" in modifications: + type_coef = float( + modifications.split("insert_typeloss_")[1].split("_coef")[0] + ) + loss_type = loss_type * type_coef + loss_type_val = loss_type.mean() + batch["loss"] = batch["loss"] + loss_type_val + + if self.logger is not None: + batch_info = self._get_batch_info(batch) + t_log = paddle.clip(batch["t"][0].clone().cast("int64"), 0, 999) + + self.logger.add( + f"loss:lattice:{mode}", loss_lat_val.item(), stack_after_epoch=True + ) + loss_lat_t = paddle.zeros([self.num_steps], dtype=loss_lat.dtype) + loss_lat_t[t_log] = loss_lat.clone() + self.logger.add( + f"loss:lattice4time:{mode}", loss_lat_t, stack_after_epoch=True + ) + + self.logger.add( + f"loss:coord:{mode}", loss_frac_val.item(), stack_after_epoch=True + ) + loss_frac_t = paddle.zeros([self.num_steps], dtype=loss_frac.dtype) + loss_frac_t[t_log] = mean_interleave( + loss_frac.clone(), batch_info["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_val.item(), + stack_after_epoch=True, + ) + loss_type_t = paddle.zeros([self.num_steps], dtype=loss_type.dtype) + loss_type_t[t_log] = mean_interleave( + loss_type.clone(), batch_info["num_atoms"] + ) + self.logger.add( + f"loss:type4time:{mode}", loss_type_t, stack_after_epoch=True + ) + + return batch + + def sampling_procedure(self, model, batch, progress_printer): + batch["xt"] = self.prior_sample(batch) + reverse_time_iterator = self.time_distribution.reverse_time_iterator( + batch, start_from=self.num_steps - 1 + ) + for t_vector in reverse_time_iterator: + batch["t"] = self.time_distribution.to_cuda(t_vector, batch) + t_value = batch["t"][0][0].item() + progress_printer(t_value) + 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): + return [ + self.lat_diffusion.output_transform(x0[0], batch), + self.frac_diffusion.output_transform(x0[1], batch), + ( + self.type_diffusion.output_transform(x0[2], batch) + if self.gen_type + else x0[2] + ), + ] + + +class DiffCSP(CrystalGen): + """DiffCSP with two-step reverse sampling.""" + + def reverse_step_sample(self, xt, t, model, batch): + lt, ft, at = xt + + # Step 1: only use fractional prediction + _, f_pred, _ = self.model_prediction(xt, t, model, batch) + ft_05 = self.frac_diffusion.reverse_step_sample_part_1(f_pred, ft, t[1], batch) + xt_05 = [lt, ft_05, at] + + # Step 2: full prediction with corrected fractional coords + l_pred, f_pred, a_pred = self.model_prediction(xt_05, t, model, batch) + lt_1 = self.lat_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) + ft_1 = self.frac_diffusion.reverse_step_sample_part_2( + f_pred, ft_05, t[1], batch + ) + 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] diff --git a/ppmat/models/miad/default_configs.py b/ppmat/models/miad/default_configs.py new file mode 100644 index 00000000..c400ba8f --- /dev/null +++ b/ppmat/models/miad/default_configs.py @@ -0,0 +1,88 @@ +# 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. + + +class _Cfg: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + if isinstance(v, dict): + setattr(self, k, _Cfg(**v)) + else: + setattr(self, k, v) + + +def _make_config(method, lat_diffusion, frac_diffusion, type_diffusion): + return _Cfg( + method=method, + cont_time=False, + num_steps=1000, + task="", + lat_diffusion=lat_diffusion, + frac_diffusion=frac_diffusion, + type_diffusion=type_diffusion, + ) + + +_lat = { + "diffcsp_ddpm_config": _Cfg(method="ddpm", scheduler="diffcsp_cosine"), + "ddpm_config": _Cfg(method="ddpm", scheduler="cosine"), + "fm_config": _Cfg(method="fm", parameterization="eps"), + "fm_v_config": _Cfg(method="fm", parameterization="v"), + "fm_lenang_config": _Cfg(method="fm_lenang"), +} + +_coord = { + "wrapped_normal_config": _Cfg( + method="wrapped_normal", scheduler="default_wrapped_normal" + ), + "pfm_config": _Cfg(method="pfm"), +} + +_type = { + "d3pm_config": _Cfg(method="d3pm", scheduler="default_d3pm"), + "ddpm_onehot_config": _Cfg(method="ddpm_onehot", scheduler="diffcsp_cosine"), +} + +DEFAULT_DIFFUSION_CONFIGS = { + "DiffCSP:diffcsp_ddpm_config:wrapped_normal_config:d3pm_config": _make_config( + "DiffCSP", + _lat["diffcsp_ddpm_config"], + _coord["wrapped_normal_config"], + _type["d3pm_config"], + ), + "Default:ddpm_config:wrapped_normal_config:d3pm_config": _make_config( + "Default", + _lat["ddpm_config"], + _coord["wrapped_normal_config"], + _type["d3pm_config"], + ), + "Default:fm_config:wrapped_normal_config:d3pm_config": _make_config( + "Default", + _lat["fm_config"], + _coord["wrapped_normal_config"], + _type["d3pm_config"], + ), + "DiffCSP:ddpm_config:wrapped_normal_config:ddpm_onehot_config": _make_config( + "DiffCSP", + _lat["ddpm_config"], + _coord["wrapped_normal_config"], + _type["ddpm_onehot_config"], + ), + "DiffCSP:fm_config:wrapped_normal_config:d3pm_config": _make_config( + "DiffCSP", + _lat["fm_config"], + _coord["wrapped_normal_config"], + _type["d3pm_config"], + ), +} diff --git a/ppmat/models/miad/diffusion_utils.py b/ppmat/models/miad/diffusion_utils.py new file mode 100644 index 00000000..db122066 --- /dev/null +++ b/ppmat/models/miad/diffusion_utils.py @@ -0,0 +1,96 @@ +# 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. + +""" +Diffusion utilities for MiAD model. +""" + +import os + +import paddle + +from paddle_scatter import scatter_mean + + +class TimeDistribution: + """Time sampling and iteration utilities.""" + + def __init__(self, num_steps, cont_time): + self.num_steps = num_steps + self.cont_time = cont_time + self.eps = 1e-3 + modifications = os.environ.get("MODIFICATIONS_FIELD", "") + if "time_log_normal_0_1_clip" in modifications: + print("MODIF: time_log_normal_0_1_clip", flush=True) + + def _atom_expand(self, batch, t): + if "batch" in batch: + num_atoms = batch["batch"].num_atoms + elif "batch_idx" in batch: + num_atoms = batch["num_atoms"] + else: + raise ValueError("Cannot determine num_atoms from batch") + + t_per_atom = t.repeat_interleave(num_atoms) + return [t, t_per_atom] + + def sample(self, batch, mode=None): + modifications = os.environ.get("MODIFICATIONS_FIELD", "") + if "time_log_normal_0_1_clip" in modifications: + s_eps = 2e-2 + t = paddle.clip( + (1 + 2 * s_eps) + * paddle.nn.functional.sigmoid(paddle.randn(batch["batch_size"])) + - s_eps, + 0, + 1, + ) + else: + t = paddle.rand([batch["batch_size"]]) + + if "max_time_" in modifications: + max_time = float(modifications.split("max_time_")[1].split("+")[0]) + t = t * max_time + + t = self.eps + (self.num_steps - 1 - self.eps) * t + if not self.cont_time: + t = t.round().cast("int64") + + return self._atom_expand(batch, t) + + def reverse_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 self._atom_expand(batch, t) + + def to_cuda(self, t_vector, batch): + return t_vector + + +def mean_interleave(t, num_repeats): + """Compute per-group mean using paddle_scatter.scatter_mean. + + Args: + t: Tensor of shape [total_atoms] with per-atom values. + num_repeats: Tensor of shape [batch_size] with atom counts per crystal. + + Returns: + Tensor of shape [batch_size] with per-crystal means. + """ + batch_idx = paddle.repeat_interleave( + paddle.arange(num_repeats.shape[0]), num_repeats + ) + return scatter_mean(t, batch_idx, dim=0) diff --git a/ppmat/models/miad/frac_diffusion.py b/ppmat/models/miad/frac_diffusion.py new file mode 100644 index 00000000..392b4f22 --- /dev/null +++ b/ppmat/models/miad/frac_diffusion.py @@ -0,0 +1,183 @@ +# 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. + +""" +Fractional coordinate diffusion models for MiAD. +""" + +import os + +import paddle + +from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler + + +class WrappedNormal: + """Wrapped Normal diffusion for fractional coordinates.""" + + def __init__(self, diffusion_config): + self.config = diffusion_config.frac_diffusion + self.num_steps = diffusion_config.num_steps + + sigmas_t, sigmas_norm_t, sb, se, d_log_p = get_scheduler( + self.config.scheduler, self.num_steps + ) + + self.sigmas_t = sigmas_t[:, None] + self.sigmas_norm_t = sigmas_norm_t[:, None] + self.sb = sb + self.d_log_p = d_log_p + + _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, + } + self.step_lr = getattr(self.config, "step_lr", None) + if self.step_lr is None: + self.step_lr = _DEFAULT_GAMMA.get(diffusion_config.task, 1e-5) + + self.drift_step_coef = paddle.ones([self.num_steps], dtype="float32")[:, None] + self.diff_step_coef = paddle.ones([self.num_steps], dtype="float32")[:, None] + self.to_domain = lambda x: x + + def output_transform(self, x0, batch): + return x0 + + def forward_step_sample(self, x0, t, batch): + st = self.sigmas_t[t.cast("int64")] + self.randn_x = paddle.randn(x0.shape) + xt = (x0 + st * self.randn_x) % 1.0 + return xt + + def reverse_step_sample_part_1(self, normed_score_pred, xt, t, batch): + 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 + std_x = paddle.sqrt(2 * step_size) + drift = ( + -step_size + * normed_score_pred + * paddle.sqrt(snt) + * self.drift_step_coef[t_idx] + ) + diffusion = std_x * paddle.randn(xt.shape) * self.diff_step_coef[t_idx] + xt_05 = xt + drift + diffusion + return xt_05 + + def reverse_step_sample_part_2(self, normed_score_pred, xt_05, t, batch): + 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 + std_x = paddle.sqrt((st_1**2 * (st**2 - st_1**2)) / (st**2)) + drift = ( + -step_size + * normed_score_pred + * paddle.sqrt(snt) + * self.drift_step_coef[t_idx] + ) + diffusion = std_x * paddle.randn(xt_05.shape) * self.diff_step_coef[t_idx] + xt_1 = (xt_05 + drift + diffusion) % 1.0 + return xt_1 + + def reverse_step_sample(self, normed_score_pred, xt, t, batch): + xt_05 = self.reverse_step_sample_part_1(normed_score_pred, xt, t, batch) + xt_1 = self.reverse_step_sample_part_2(normed_score_pred, xt_05, t, batch) + return xt_1 + + def prior_sample(self, batch): + num_atoms = batch["num_atoms"] + total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + return paddle.rand([total_atoms, 3], dtype="float32") + + def loss(self, batch): + t_idx = batch["t"][1].cast("int64") + st = self.sigmas_t[t_idx] + snt = self.sigmas_norm_t[t_idx] + normed_score_pred = batch["prediction"][1] + normed_score = self.d_log_p(st * self.randn_x, st) / paddle.sqrt(snt) + l2 = ( + ((normed_score_pred - normed_score) ** 2) + .reshape([normed_score_pred.shape[0], -1]) + .mean(axis=1) + ) + + # mirage infusion code + modifications = os.environ.get("MODIFICATIONS_FIELD", "") + if "miad:add_mirage_atoms_upto" in modifications: + mirage_type = 0 + atom_types = batch["x0"][2] + mask = (atom_types != mirage_type).cast(l2.dtype) + coef = mask.shape[0] / mask.sum() + l2 = l2 * mask * coef + + return l2 + + def get_x0_prediction(self, normed_score_pred, xt, t, batch): + t_idx = t.cast("int64") + x0_pred = (xt - self.sigmas_t[t_idx] * normed_score_pred) % 1.0 + return x0_pred + + +class PFM: + """Periodic Flow Matching for fractional coordinates.""" + + def __init__(self, diffusion_config): + self.config = diffusion_config.frac_diffusion + self.num_steps = diffusion_config.num_steps + self.cont_time = diffusion_config.cont_time + self.final_t = diffusion_config.num_steps + self.step_coef = paddle.ones([self.final_t], dtype="float32")[:, None] + self.step = paddle.to_tensor(1.0 / self.final_t) + self.to_domain = lambda x: x + self.default_loss_scale = 10 + + def output_transform(self, x0, batch): + return x0 + + def forward_step_sample(self, x0, t, batch): + xT_minus_x0 = paddle.rand(x0.shape) - 0.5 + step = (1 + t[:, None]) * self.step + xt = (x0 + step * xT_minus_x0) % 1.0 + self.ut = (-1) * xT_minus_x0 + return xt + + def reverse_step_sample(self, vt, xt, t, batch): + t_idx = t.cast("int64") + xt_1 = (xt + self.step_coef[t_idx] * self.step * vt) % 1.0 + return xt_1 + + def prior_sample(self, batch): + num_atoms = batch["num_atoms"] + total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + return paddle.rand([total_atoms, 3]) + + def loss(self, batch): + vt = batch["prediction"][1] + l2 = ((vt - self.ut) ** 2).reshape([-1, 3]).mean(axis=1) + return l2 * self.default_loss_scale + + def get_x0_prediction(self, vt, xt, t, batch): + step = (1 + t[:, None]) * self.step + return (xt + step * vt) % 1.0 diff --git a/ppmat/models/miad/graph_utils.py b/ppmat/models/miad/graph_utils.py new file mode 100644 index 00000000..150618e9 --- /dev/null +++ b/ppmat/models/miad/graph_utils.py @@ -0,0 +1,52 @@ +# 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. + +"""Graph utility for converting dense adjacency to COO edge_index.""" + +import paddle + + +def dense_to_sparse(adj): + """Convert a dense adjacency matrix to sparse (edge_index, edge_attr) COO format. + + Supports both single-graph (2D) and batched (3D) dense adjacency tensors. + """ + if adj.ndim == 2: + nonzero = paddle.nonzero(adj.cast("float32")) + if nonzero.shape[0] == 0: + return paddle.zeros([2, 0], dtype="int64"), paddle.zeros( + [0], dtype="float32" + ) + edge_index = nonzero.t() + return edge_index, paddle.ones([edge_index.shape[1]], dtype="float32") + + # Batched case + batch_size, num_nodes, _ = adj.shape + edge_indices_list = [] + edge_attrs_list = [] + + for b in range(batch_size): + adj_b = adj[b] + nonzero = paddle.nonzero(adj_b.cast("float32")) + if nonzero.shape[0] == 0: + continue + edge_indices_list.append(nonzero.t()) + edge_attrs_list.append(paddle.ones([nonzero.shape[0]], dtype="float32")) + + if not edge_indices_list: + return paddle.zeros([2, 0], dtype="int64"), paddle.zeros([0], dtype="float32") + + edge_index = paddle.concat(edge_indices_list, axis=1) + edge_attr = paddle.concat(edge_attrs_list, axis=0) + return edge_index, edge_attr diff --git a/ppmat/models/miad/lattice_diffusion.py b/ppmat/models/miad/lattice_diffusion.py new file mode 100644 index 00000000..f0af3670 --- /dev/null +++ b/ppmat/models/miad/lattice_diffusion.py @@ -0,0 +1,253 @@ +# 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. + +""" +Lattice diffusion models for MiAD. +""" + +import paddle + +from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler + + +class DDPM: + """DDPM for lattice diffusion. + + Note: This is a lightweight DDPM that pre-computes and stores coefficient + tensors for direct indexing during forward/reverse steps. It differs from + ppmat.schedulers.scheduling_ddpm.DDPMScheduler which is a full-featured + class-based scheduler. The direct-index design is required by MiAD's + CrystalGen which indexes coeff tensors by timestep in tight loops, where + the class-based interface overhead would be prohibitive. + """ + + def __init__(self, diffusion_config, scheduler_cfg=None): + self.config = ( + scheduler_cfg + if scheduler_cfg is not None + else diffusion_config.lat_diffusion + ) + self.num_steps = diffusion_config.num_steps + self.cont_time = diffusion_config.cont_time + + cumprod_alphas_t, ca_t = get_scheduler(self.config.scheduler, self.num_steps) + + cumprod_alphas_t_1 = cumprod_alphas_t[:-1] + cumprod_alphas_t = cumprod_alphas_t[1:] + + if ca_t is not None: + self.f_cumprod_alphas_t = lambda t: ca_t(1 + t).reshape([-1, 1, 1]) + else: + self.f_cumprod_alphas_t = None + + alphas_t = cumprod_alphas_t / cumprod_alphas_t_1 + betas_t = 1 - alphas_t + + # Reshape to (num_steps, 1, 1) for broadcasting with (batch, 3, 3) + self.betas_t = betas_t.reshape([-1, 1, 1]) + self.alphas_t = alphas_t.reshape([-1, 1, 1]) + self.cumprod_alphas_t = cumprod_alphas_t.reshape([-1, 1, 1]) + self.cumprod_alphas_t_1 = cumprod_alphas_t_1.reshape([-1, 1, 1]) + + # Pre-compute reverse sampling coefficients + self.reverse_c0 = 1 / paddle.sqrt(self.alphas_t) + self.reverse_c1 = (1 - self.alphas_t) / paddle.sqrt(1 - self.cumprod_alphas_t) + self.reverse_std_coef = paddle.sqrt( + self.betas_t * (1 - self.cumprod_alphas_t_1) / (1 - self.cumprod_alphas_t) + ) + + # Pre-compute eps-to-x0 coefficients + self.eps_to_x0_c0 = paddle.sqrt(1 / self.cumprod_alphas_t) + self.eps_to_x0_c1 = paddle.sqrt(1 / self.cumprod_alphas_t - 1) + + self.to_domain = lambda x: x + + def output_transform(self, x0, batch): + return x0 + + def forward_step_sample(self, x0, t, batch): + if self.cont_time and self.f_cumprod_alphas_t is not None: + at = self.f_cumprod_alphas_t(t) + else: + at = self.cumprod_alphas_t[t.cast("int64")] + self.randn_x = paddle.randn(x0.shape) + xt = paddle.sqrt(at) * x0 + paddle.sqrt(1 - at) * self.randn_x + return xt + + def reverse_step_sample(self, eps_pred, xt, t, batch): + t_idx = t.cast("int64") + mu_xt_1 = self.reverse_c0[t_idx] * (xt - self.reverse_c1[t_idx] * eps_pred) + if (t_idx == 0).all(): + return mu_xt_1 + return mu_xt_1 + self.reverse_std_coef[t_idx] * paddle.randn(mu_xt_1.shape) + + def prior_sample(self, batch): + return paddle.randn([batch["batch_size"], 3, 3], dtype="float32") + + def loss(self, batch): + eps_pred = batch["prediction"][0] + l2 = ( + ((eps_pred - self.randn_x) ** 2) + .reshape([eps_pred.shape[0], -1]) + .mean(axis=1) + ) + return l2 + + def get_x0_prediction(self, eps_pred, xt, t, batch): + t_idx = t.cast("int64") + x0_pred = self.eps_to_x0_c0[t_idx] * xt - self.eps_to_x0_c1[t_idx] * eps_pred + return x0_pred + + +class FM: + """Flow Matching for lattice diffusion.""" + + def __init__(self, diffusion_config): + self.config = diffusion_config.lat_diffusion + self.num_steps = diffusion_config.num_steps + self.cont_time = diffusion_config.cont_time + self.step = 1.0 / self.num_steps + self.parameterization = getattr( + diffusion_config.lat_diffusion, "parameterization", "eps" + ) + self.step_coef = paddle.ones([self.num_steps], dtype="float32")[:, None, None] + self.to_domain = lambda x: x + + def output_transform(self, x0, batch): + return x0 + + def forward_step_sample(self, x0, t, batch): + eps = paddle.randn(x0.shape) + step = (1 + t[:, None, None]) * self.step + xt = (1 - step) * x0 + step * eps + if self.parameterization == "eps": + self.ut = (-1) * eps + elif self.parameterization == "v": + self.ut = x0 - eps + return xt + + def reverse_step_sample(self, pred, xt, t, batch): + t_idx = t.cast("int64") + if self.parameterization == "eps": + if t_idx[0] >= 999: + return paddle.randn(xt.shape) + eps_pred = (-1) * pred + step = (1 + t[:, None, None]) * self.step + x0_pred = (xt - step * eps_pred) / (1 - step) + vt = x0_pred - eps_pred + elif self.parameterization == "v": + vt = pred + xt_1 = xt + self.step_coef[t_idx] * self.step * vt + return xt_1 + + def prior_sample(self, batch): + return paddle.randn([batch["batch_size"], 3, 3], dtype="float32") + + def loss(self, batch): + vt = batch["prediction"][0] + l2 = ((vt - self.ut) ** 2).reshape([-1, 9]).mean(axis=1) + return l2 + + def get_x0_prediction(self, pred, xt, t, batch): + step = (1 + t[:, None, None]) * self.step + t_idx = t.cast("int64") + if self.parameterization == "eps": + if t_idx[0] >= 999: + return paddle.randn(xt.shape) + eps_pred = (-1) * pred + x0_pred = (xt - step * eps_pred) / (1 - step) + vt = x0_pred - eps_pred + elif self.parameterization == "v": + vt = pred + return xt + step * vt + + +class FM_LenAng: + """Flow Matching for lattice in length-angle parameterization.""" + + def __init__(self, diffusion_config): + self.config = diffusion_config.lat_diffusion + self.num_steps = diffusion_config.num_steps + self.cont_time = diffusion_config.cont_time + self.step = 1.0 / self.num_steps + + alpha, theta = 1.3, 0.25 + self.gamma_alpha = alpha + self.gamma_theta = theta + self.angle_difference_bound = 20 + self.default_loss_scale = 0.45 + + self._lat2lenang = None + self._lenang2lat = None + self.to_domain = lambda x: x + + def _get_converters(self): + if self._lat2lenang is None: + from ppmat.utils.crystal import lattice_params_to_matrix_paddle + from ppmat.utils.crystal import lattices_to_params_shape_paddle + + self._lat2lenang = lambda lat: paddle.concat( + lattices_to_params_shape_paddle(lat), axis=1 + ) + self._lenang2lat = lambda la: lattice_params_to_matrix_paddle( + la[:, :3], la[:, 3:] + ) + return self._lat2lenang, self._lenang2lat + + def output_transform(self, x0, batch): + return x0 + + def forward_step_sample(self, x0, t, batch): + xT = self.prior_sample(batch) + step = (1 + t[:, None, None]) * self.step + xt = (1 - step) * x0 + step * xT + self.ut = x0 - xT + return xt + + def reverse_step_sample(self, vt, xt, t, batch): + xt_1 = xt + self.step * vt + return xt_1 + + def prior_sample(self, batch): + bs = batch["batch_size"] + lenang_xT = paddle.zeros([bs, 6], dtype="float32") + + gamma_dist = paddle.distribution.Gamma( + paddle.to_tensor([self.gamma_alpha], dtype="float32"), + paddle.to_tensor([self.gamma_theta], dtype="float32"), + ) + lenang_xT[:, :3] = 2 + gamma_dist.sample([bs, 3]).reshape([bs, 3]) + + # Sample angles from constrained uniform + ang = 60 + 60 * paddle.rand([4 * 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) + ) + valid = ang[check][:bs] + lenang_xT[:, 3:] = valid + + _, lenang2lat = self._get_converters() + xT = lenang2lat(lenang_xT) + return xT + + def loss(self, batch): + vt = batch["prediction"][0] + l2 = ((vt - self.ut) ** 2).reshape([-1, 9]).mean(axis=1) + return self.default_loss_scale * l2 + + def get_x0_prediction(self, vt, xt, t, batch): + step = (1 + t[:, None, None]) * self.step + return xt + step * vt diff --git a/ppmat/models/miad/miad.py b/ppmat/models/miad/miad.py new file mode 100644 index 00000000..9f21c94d --- /dev/null +++ b/ppmat/models/miad/miad.py @@ -0,0 +1,256 @@ +# 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 types import SimpleNamespace + +import numpy as np +import paddle +import paddle.nn as nn + +from ppmat.models.miad.crystal_diffusion import init_diffusion +from ppmat.models.miad.crystal_diffusion import parse_num_atoms_to_per_crystal +from ppmat.models.miad.miad_cspnet import CSPNet +from ppmat.utils import logger +from ppmat.utils.crystal import lattices_to_params_shape_numpy + + +def _dict_to_sns(d): + """Recursively convert dict to SimpleNamespace for attribute access.""" + if not isinstance(d, dict): + return d + return SimpleNamespace(**{k: _dict_to_sns(v) for k, v in d.items()}) + + +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 {} + + self.decoder = CSPNet(**model_cfg) + if isinstance(diffusion_cfg, dict): + diffusion_cfg = _dict_to_sns(diffusion_cfg) + self.diffusion = init_diffusion(diffusion_cfg, logger=None) + + def set_state_dict(self, state_dict, use_structured_name=True): + """ + Load checkpoint with automatic legacy format detection. + + MiAD wraps CSPNet as self.decoder, so all parameter keys in a native + Paddle checkpoint are prefixed with "decoder." + (e.g. decoder.csp_layer_0.edge_mlp.0.weight). + + Legacy PyTorch checkpoints have bare keys (e.g. csp_layer_0.edge_mlp.0.weight) + and Linear weights stored in (out_features, in_features) format. + This method auto-detects and converts them: adds decoder. prefix and transposes + all 2D weight tensors to Paddle's (in_features, out_features) format. + + For legacy checkpoints, ALL 2D Linear weights are unconditionally transposed. + Shape-based detection cannot distinguish square matrices (e.g. 512x512) where + both (out,in) and (in,out) have the same shape. The only safe approach is to + always transpose for legacy-format checkpoints. + """ + model_state = self.state_dict() + model_keys = set(model_state.keys()) + state_keys = set(state_dict.keys()) + + if len(state_keys) == 0: + return super().set_state_dict(state_dict, use_structured_name) + + has_decoder_prefix = any(k.startswith("decoder.") for k in state_keys) + keys_native = state_keys == model_keys or state_keys.issubset(model_keys) + + if keys_native and has_decoder_prefix: + return super().set_state_dict(state_dict, use_structured_name) + + is_legacy = not has_decoder_prefix + processed = {} + num_transposed = 0 + num_copied = 0 + for k, v in state_dict.items(): + new_key = f"decoder.{k}" if is_legacy else k + val_np = v.numpy() if hasattr(v, "numpy") else v + + # Legacy checkpoints always store Linear weights in PyTorch format + # (out_features, in_features). For square matrices (e.g. 512x512), + # shape-based detection cannot work because both orientations have + # identical shapes. Therefore we unconditionally transpose all 2D + # Linear weights when loading from a legacy checkpoint. + if is_legacy and new_key.endswith(".weight") and len(v.shape) == 2: + val_np = val_np.T + num_transposed += 1 + elif is_legacy: + num_copied += 1 + + processed[new_key] = val_np + + model_in_keys = set(model_keys) - set(processed.keys()) + extra_in_ckpt = set(processed.keys()) - set(model_keys) + if model_in_keys: + logger.info( + f"[MiAD] {len(model_in_keys)} model keys missing from checkpoint" + f" (buffers/renamed): {sorted(model_in_keys)[:3]}..." + ) + if extra_in_ckpt: + logger.info( + f"[MiAD] {len(extra_in_ckpt)} checkpoint keys not in model" + f" (skipped): {sorted(extra_in_ckpt)[:3]}..." + ) + if is_legacy: + logger.info( + f"[MiAD] Added decoder. prefix; transposed {num_transposed}" + f" Linear weights, direct copied {num_copied} non-Linear params" + ) + + # Load parameters one by one via set_value to bypass Paddle's + # set_state_dict that may report success without actually loading values. + loaded_missing = [] + loaded_unexpected = [] + for name, param in self.named_parameters(): + if name in processed: + proc_val = processed[name] + if proc_val.shape != tuple(param.shape): + logger.warning( + f"[MiAD] SHAPE MISMATCH: {name}, model={param.shape}," + f" loaded={proc_val.shape} - SKIPPED" + ) + continue + param.set_value(proc_val) + else: + loaded_missing.append(name) + + extra_loaded = set(processed.keys()) - set(model_keys) + if extra_loaded: + loaded_unexpected = list(extra_loaded) + + loaded = (loaded_missing, loaded_unexpected) + return loaded + + def forward(self, batch, **kwargs): + mode = "train" if self.training else "val" + 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 = { + "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 + + batch = self.diffusion.sampling_procedure( + model=self.decoder, + batch=batch_data, + progress_printer=progress_printer, + ) + + 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 = lat_i.numpy() if hasattr(lat_i, "numpy") else np.asarray(lat_i) + fc_np = fc_i.numpy() if hasattr(fc_i, "numpy") else np.asarray(fc_i) + at_np = at_i.numpy() if hasattr(at_i, "numpy") else np.asarray(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/miad_cspnet.py b/ppmat/models/miad/miad_cspnet.py new file mode 100644 index 00000000..d5b9a278 --- /dev/null +++ b/ppmat/models/miad/miad_cspnet.py @@ -0,0 +1,133 @@ +# 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. + +""" +CSPNet for MiAD model. Inherits core structure from ppmat.models.diffcsp.diffcsp.CSPNet. +Overrides gen_edges for block_diag edge construction and forward for pre-computed +time embeddings and flexible atom type (one-hot / discrete) handling. +""" + +import paddle + +from paddle_scatter import scatter +from ppmat.models.diffcsp.diffcsp import CSPNet as _DiffCSPCSPNet +from ppmat.models.miad.graph_utils import dense_to_sparse + + +class CSPNet(_DiffCSPCSPNet): + """MiAD-specific CSPNet inheriting from DiffCSP CSPNet. + + Key differences from the parent: + - gen_edges: uses paddle.block_diag + dense_to_sparse for fully-connected edges + - forward: accepts pre-computed t_emb (time embedding) and handles both + one-hot and discrete atom types + - Removes prop_mlp from each CSPLayer (MiAD does not use property embeddings) + """ + + def __init__( + self, + hidden_dim=128, + latent_dim=256, + num_layers=4, + max_atoms=100, + act_fn="silu", + dis_emb="sin", + num_freqs=10, + edge_style="fc", + ln=False, + ip=True, + smooth=False, + pred_type=False, + cutoff=7.0, + max_neighbors=20, + model_name=None, + **kwargs, + ): + self.cutoff = cutoff + self.max_neighbors = max_neighbors + super().__init__( + hidden_dim=hidden_dim, + latent_dim=latent_dim, + num_layers=num_layers, + act_fn=act_fn, + dis_emb=dis_emb, + num_freqs=num_freqs, + edge_style=edge_style, + ln=ln, + ip=ip, + smooth=smooth, + pred_type=pred_type, + num_classes=max_atoms, + ) + for i in range(self.num_layers): + layer = self._modules.get(f"csp_layer_{i}") + if layer and hasattr(layer, 'prop_mlp'): + del layer.prop_mlp + + def gen_edges(self, num_atoms, frac_coords): + """Fully-connected graph via block diagonal adjacency matrix.""" + if self.edge_style == "fc": + lis = [paddle.ones([n, n], dtype=num_atoms.dtype) for n in num_atoms] + fc_graph = paddle.block_diag(lis) + fc_edges, _ = dense_to_sparse(fc_graph) + return fc_edges, (frac_coords[fc_edges[1]] - frac_coords[fc_edges[0]]) + + def forward( + self, t, t_emb, atom_types, frac_coords, lattices, num_atoms, node2graph + ): + edges, frac_diff = self.gen_edges(num_atoms, frac_coords) + edge2graph = node2graph[edges[0]] + # Handle both discrete atom types and one-hot encoding + if atom_types.ndim > 1: + if self.smooth: + node_features = self.node_embedding(atom_types.cast("float32")) + else: + atom_indices = atom_types.argmax(axis=-1) + node_features = self.node_embedding(atom_indices.cast("int64")) + else: + if self.smooth: + node_features = self.node_embedding(atom_types.cast("float32")) + else: + node_features = self.node_embedding(atom_types - 1) + + t_per_atom = paddle.repeat_interleave(t_emb, num_atoms, axis=0) + node_features = paddle.concat([node_features, t_per_atom], axis=1) + node_features = self.atom_latent_emb(node_features) + + for i in range(0, self.num_layers): + node_features = self._modules["csp_layer_%d" % i]( + node_features, + frac_coords, + lattices, + edges, + edge2graph, + frac_diff=frac_diff, + ) + + if self.ln: + node_features = self.final_layer_norm(node_features) + + coord_out = self.coord_out(node_features) + + graph_features = scatter(node_features, node2graph, dim=0, reduce="mean") + lattice_out = self.lattice_out(graph_features) + lattice_out = lattice_out.reshape([-1, 3, 3]) + + if self.ip: + lattice_out = paddle.einsum("bij,bjk->bik", lattice_out, lattices) + if self.pred_type: + type_out = self.type_out(node_features) + return lattice_out, coord_out, type_out + + return lattice_out, coord_out diff --git a/ppmat/models/miad/type_diffusion.py b/ppmat/models/miad/type_diffusion.py new file mode 100644 index 00000000..74447cc3 --- /dev/null +++ b/ppmat/models/miad/type_diffusion.py @@ -0,0 +1,222 @@ +# 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. + +""" +Atom type diffusion models for MiAD. +""" + +import paddle +import paddle.nn.functional as F + +from ppmat.models.miad.lattice_diffusion import DDPM +from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler + + +class DDPM_onehot(DDPM): + """DDPM for atom types with one-hot encoding.""" + + def __init__(self, diffusion_config): + super().__init__( + diffusion_config, scheduler_cfg=diffusion_config.type_diffusion + ) + + # Reshape coefficients from (N,1,1) to (N,1) for atom types + _reshape_attrs = [ + "betas_t", + "alphas_t", + "cumprod_alphas_t", + "cumprod_alphas_t_1", + "reverse_c0", + "reverse_c1", + "reverse_std_coef", + "eps_to_x0_c0", + "eps_to_x0_c1", + ] + for attr in _reshape_attrs: + value = getattr(self, attr, None) + if value is not None and isinstance(value, paddle.Tensor): + if len(value.shape) == 3: + setattr(self, attr, value.reshape([-1, 1])) + + self.num_types = 100 + 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 output_transform(self, x0, batch): + return x0 + + def forward_step_sample(self, x0, t, batch): + onehot_x0 = self.to_domain(x0) + onehot_xt = super().forward_step_sample(onehot_x0, t, batch) + return onehot_xt + + def reverse_step_sample(self, onehot_eps_pred, onehot_xt, t, batch): + onehot_xt_1 = super().reverse_step_sample(onehot_eps_pred, onehot_xt, t, batch) + if (t.cast("int64") == 0).all(): + xt_1 = self.from_domain(onehot_xt_1) + return xt_1 + return onehot_xt_1 + + def prior_sample(self, batch): + num_atoms = batch["num_atoms"] + total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + return paddle.randn([total_atoms, self.num_types], dtype="float32") + + def loss(self, batch): + eps_pred = batch["prediction"][2] + l2 = ( + ((eps_pred - self.randn_x) ** 2) + .reshape([eps_pred.shape[0], -1]) + .mean(axis=1) + ) + return l2 + + def get_x0_prediction(self, onehot_eps_pred, onehot_xt, t, batch, x0_format="disc"): + onehot_x0_pred = super().get_x0_prediction(onehot_eps_pred, onehot_xt, t, batch) + if x0_format == "onehot": + return onehot_x0_pred + elif x0_format == "disc": + return self.from_domain(onehot_x0_pred) + + +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.num_types = 100 + + Q_t, cumprod_Q_t = get_scheduler(self.config.scheduler, self.num_steps) + + Q_t_1 = paddle.concat( + [ + paddle.eye(Q_t.shape[1], dtype="float32").unsqueeze(0), + Q_t[:-1], + ], + axis=0, + ) + self.Q_t = Q_t.reshape([-1, self.num_types, self.num_types]) + self.Q_t_1 = Q_t_1.reshape([-1, self.num_types, self.num_types]) + cumprod_Q_t_1 = paddle.concat( + [ + paddle.eye(cumprod_Q_t.shape[1], dtype="float32").unsqueeze(0), + cumprod_Q_t[:-1], + ], + axis=0, + ) + self.cumprod_Q_t = cumprod_Q_t.reshape([-1, self.num_types, self.num_types]) + self.cumprod_Q_t_1 = cumprod_Q_t_1.reshape([-1, self.num_types, self.num_types]) + + 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") + ) + onehot_xt = self.to_domain(xt) + return onehot_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) + result = paddle.where( + paddle.isnan(result), + paddle.full_like(result, 1.0 / self.num_types), + result, + ) + return 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) + xt_1_probs = paddle.where( + paddle.isnan(xt_1_probs), + paddle.full_like(xt_1_probs, 1.0 / self.num_types), + xt_1_probs, + ) + 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") + ) + onehot_xt_1 = self.to_domain(xt_1) + return onehot_xt_1 + + def prior_sample(self, batch): + num_atoms = batch["num_atoms"] + total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + shape = [total_atoms, 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") + ) + onehot_xT = self.to_domain(xT) + return onehot_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-4 + 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 + + def get_x0_prediction(self, onehot_pred, onehot_xt, t, batch): + onehot_x0 = self.prediction_to_domain(onehot_pred) + return self.from_domain(onehot_x0) diff --git a/ppmat/schedulers/__init__.py b/ppmat/schedulers/__init__.py index 8cd773f1..485c0c84 100644 --- a/ppmat/schedulers/__init__.py +++ b/ppmat/schedulers/__init__.py @@ -22,12 +22,16 @@ from ppmat.schedulers.scheduling_lattice_vp import LatticeVPSDEScheduler from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeScheduler from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeSchedulerWrapped +from ppmat.schedulers.miad_schedulers import scheduler as build_miad_scheduler +from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal +from ppmat.schedulers.scheduling_sde_ve import p_wrapped_normal NumAtomsVarianceAdjustedWrappedVESDE = ( scheduling_wrapped_sde_ve.NumAtomsVarianceAdjustedWrappedVESDE ) __all__ = [ "build_scheduler", + "build_miad_scheduler", "DDPMScheduler", "ScoreSdeVeScheduler", "ScoreSdeVeSchedulerWrapped", @@ -35,6 +39,8 @@ "NumAtomsVarianceAdjustedWrappedVESDE", "D3PMScheduler", "NoiseScheduler", + "p_wrapped_normal", + "d_log_p_wrapped_normal", ] diff --git a/ppmat/schedulers/miad_schedulers.py b/ppmat/schedulers/miad_schedulers.py new file mode 100644 index 00000000..3f1c5e53 --- /dev/null +++ b/ppmat/schedulers/miad_schedulers.py @@ -0,0 +1,146 @@ +# 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. + +""" +Pre-computed coefficient schedulers for MiAD-style diffusion models. + +These functions return raw coefficient tensors that diffusion modules index into +during forward/reverse steps. This is a lighter-weight alternative to the class-based +scheduler interface (DDPMScheduler, ScoreSdeVeScheduler, etc.) used by DiffCSP. +""" + +import math + +import numpy as np +import paddle + +from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal +from ppmat.schedulers.scheduling_sde_ve import p_wrapped_normal # noqa: F401 + + +def scheduler(scheduler_name, num_steps): + """Factory for pre-computed diffusion coefficient tensors. + + Args: + scheduler_name: One of 'diffcsp_cosine', 'cosine', 'default_d3pm', + 'default_wrapped_normal'. + num_steps: Number of diffusion steps. + + Returns: + Tuple of coefficient tensors. The format depends on the scheduler type: + - diffcsp_cosine: (cumprod_alphas_t, None) + - cosine: (cumprod_alphas_t, a_t_fn) + - default_d3pm: (Q_t, cumprod_Q_t) + - default_wrapped_normal: (sigmas, sigmas_norm, sigma_begin, + sigma_end, d_log_p_fn) + """ + if scheduler_name == "diffcsp_cosine": + return _scheduler_diffcsp_cosine(num_steps) + elif scheduler_name == "cosine": + return _scheduler_cosine(num_steps) + elif "default_d3pm" in scheduler_name: + return _scheduler_default_d3pm(num_steps) + elif scheduler_name == "default_wrapped_normal": + return _scheduler_default_wrapped_normal(num_steps) + + raise NotImplementedError(f"Scheduler '{scheduler_name}' not implemented") + + +def _scheduler_diffcsp_cosine(num_steps): + s = 0.008 + discretization = paddle.linspace(0, num_steps, num_steps + 1) + alphas_cumprod = ( + paddle.cos(((discretization / num_steps) + s) / (1 + s) * math.pi * 0.5) ** 2 + ) + alphas_cumprod = alphas_cumprod / alphas_cumprod[0] + betas = 1 - alphas_cumprod[1:] / alphas_cumprod[:-1] + betas = paddle.clip(betas, 0.0001, 0.9999) + alphas = 1.0 - betas + cumprod_alphas_t = paddle.cumprod(alphas, dim=0).cast("float32") + cumprod_alphas_t = paddle.concat([paddle.to_tensor([1.0]), cumprod_alphas_t]) + return cumprod_alphas_t, None + + +def _scheduler_cosine(num_steps): + s = 0.008 + + def f_t(t): + return paddle.cos((t / (num_steps + 1) + s) / (1 + s) * math.pi / 2) ** 2 + + def a_t(t): + return f_t(t) / f_t(paddle.to_tensor(0.0, dtype="float64")) + + discretization = paddle.linspace(0, num_steps, num_steps + 1) + cumprod_alphas_t = a_t(discretization).cast("float32") + return cumprod_alphas_t, a_t + + +def _scheduler_default_d3pm(num_steps, vocab_size=100): + def get_uniform_transition_mat(vocab_size, beta_t): + mat = paddle.full((vocab_size, vocab_size), beta_t / float(vocab_size)) + diag_val = 1 - beta_t * (vocab_size - 1) / vocab_size + for i in range(vocab_size): + mat[i, i] = diag_val.cast("float32") + return mat + + s = 0.008 + + def f_t(t): + return paddle.cos((t / (num_steps + 1) + s) / (1 + s) * math.pi / 2) + + def a_t(t): + return f_t(t) / f_t(paddle.to_tensor(0.0, dtype="float64")) + + discretization = paddle.arange(1, num_steps + 1, dtype="float64") + cumprod_alphas_t = a_t(discretization) + 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_steps): + Q_t_list.append(get_uniform_transition_mat(vocab_size, betas_t[t_idx])) + Q_t = paddle.stack(Q_t_list, axis=0) + + cumprod_Q_t_list = [Q_t[0]] + for t_idx in range(1, num_steps): + 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) + + return Q_t, cumprod_Q_t + + +def _scheduler_default_wrapped_normal(num_steps): + """Wrapped normal scheduler for fractional coordinate diffusion. + + Uses sn=100000 for sigma_norm calculation (differs from scheduling_sde_ve.py + which uses sn=10000). This larger sample size was used in the original MiAD + implementation for more precise expected score norm estimation. + """ + sigma_begin, sigma_end = 0.005, 0.5 + sigmas = paddle.to_tensor( + np.exp(np.linspace(np.log(sigma_begin), np.log(sigma_end), num_steps)), + dtype="float32", + ) + + def sigma_norm(sigma, T=1.0, sn=100000): + sigmas_expanded = paddle.tile(sigma[None, :], [sn, 1]) + x_sample = sigma * paddle.randn(sigmas_expanded.shape) + x_sample = x_sample % T + normal_ = d_log_p_wrapped_normal(x_sample, sigmas_expanded, T=T) + return (normal_**2).mean(axis=0) + + sigmas_norm_ = sigma_norm(sigmas) + return sigmas, sigmas_norm_, sigma_begin, sigma_end, d_log_p_wrapped_normal diff --git a/structure_generation/configs/miad/README.md b/structure_generation/configs/miad/README.md new file mode 100644 index 00000000..1bd0938e --- /dev/null +++ b/structure_generation/configs/miad/README.md @@ -0,0 +1,77 @@ +# MiAD + +MiAD (Mirage Atom Diffusion) is a diffusion-based framework for de novo crystal generation. It introduces the concept of 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. + +## How It Works + +1. Mirage Atoms: The model uses sentinel nodes called "mirage atoms" (assigned atom_type = 0) that can be transformed into real chemical elements or discarded as empty space during the diffusion trajectory . + +2. Training Phase: Real crystals are padded with mirage atoms up to a maximum limit (e.g., 25 atoms), where mirage nodes have random fractional coordinates and type 0 . + +3. Generation Phase: The model initializes all crystals with the maximum atom count and predicts which nodes should materialize into real atoms versus remaining as mirage atoms . + +4. Post-Processing: Remaining mirage atoms are stripped before exporting to final CIF files . + +## Architecture + +MiAD uses DiffCSP as its backbone architecture, with the CrystalGen orchestrator handling the diffusion process for lattice parameters, fractional coordinates, and atom types. + +## Dataset + +[Original dataset address](https://drive.google.com/file/d/1BLI3VtvzfIIXlH6UHQ4o-gQaCIOZ1UR7/view?usp=sharing) + +[Mirror AIStudio address](https://aistudio.baidu.com/modelsdetail/48578/intro) + +Extract the data to `./data/mp_20/` so that the CSV files are at `./data/mp_20/train.csv`, `./data/mp_20/val.csv`, and `./data/mp_20/test.csv`. + +## Environment Dependencies + +- Python >= 3.10 +- PaddlePaddle >= 3.3.0 (official release) +- paddle_scatter >= 2.1.2 +- numpy, pymatgen, omegaconf + +## Checkpoints + +[Original checkpoints address](https://drive.google.com/file/d/1KyD6KzvjYFPfU8lutFyO_0b8EbeHSGqf/view?usp=sharing) + +[Mirror AIStudio address](https://aistudio.baidu.com/modelsdetail/48578/intro) + +[Paddle checkpoints address](https://aistudio.baidu.com/modelsdetail/48638/intro) + + + +## Commands + +### Training + +```bash +# single GPU +python structure_generation/train.py -c structure_generation/configs/miad/miad_mp20.yaml + +# multi-GPU (example: 4 GPUs) +python -m paddle.distributed.launch --gpus="0,1,2,3" structure_generation/train.py -c structure_generation/configs/miad/miad_mp20.yaml +``` + +### Validation + +```bash +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' +``` + +### Sampling (Structure Generation) + +```bash +# Option 1: Use pre-trained model (auto-download) +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 + +# Option 2: Custom checkpoint +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' +``` + +### Evaluation (Compute Metrics) + +```bash +# Evaluate generated structures against ground truth using CSPMetric +python structure_generation/sample.py --model_name='miad_mp20' --weights_name='latest.pdparams' --save_path='result_eval/' --mode='compute_metric' +``` \ No newline at end of file diff --git a/structure_generation/configs/miad/miad_mp20.yaml b/structure_generation/configs/miad/miad_mp20.yaml new file mode 100644 index 00000000..246c5ec8 --- /dev/null +++ b/structure_generation/configs/miad/miad_mp20.yaml @@ -0,0 +1,154 @@ +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: diffcsp_cosine + frac_diffusion: + method: wrapped_normal + scheduler: default_wrapped_normal + type_diffusion: + method: d3pm + scheduler: default_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: MiADCollator + 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: MiADCollator + 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: MiADCollator + +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: MiADCollator + 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/test_miad.py b/test/test_miad.py new file mode 100644 index 00000000..23093ed5 --- /dev/null +++ b/test/test_miad.py @@ -0,0 +1,114 @@ +# 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.mp20_dataset import MP20Dataset +from ppmat.models import build_model +from ppmat.models.miad.collate import MiADCollator # noqa: F401 + + +class MiADConfigTest(unittest.TestCase): + """Test that MiAD config can be parsed and model constructed.""" + + 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_model_construction(self): + config = OmegaConf.load("structure_generation/configs/miad/miad_mp20.yaml") + config = OmegaConf.to_container(config, resolve=True) + model = build_model(config["Model"]) + self.assertIsNotNone(model) + self.assertIsInstance(model, paddle.nn.Layer) + + def test_model_forward(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_size = 2 + num_atoms = paddle.to_tensor([5, 7], dtype="int64") + total_atoms = int(num_atoms.sum()) + batch_idx = paddle.concat( + [paddle.full([int(n)], i, dtype="int64") for i, n in enumerate(num_atoms)] + ) + 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") + + batch = { + "x0": [lattices, frac_coords, atom_types], + "batch_size": batch_size, + "num_atoms": num_atoms, + "batch_idx": batch_idx, + "atom_types": atom_types, + } + + with paddle.no_grad(): + output = model(batch) + + self.assertIn("loss_dict", output) + self.assertIn("loss", output["loss_dict"]) + + +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 = MiADCollator() + samples = [self.dataset[i] for i in range(min(4, len(self.dataset)))] + batch = collator(samples) + 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() From 656aa23232fbc2b44fd7855d046506416be24ded Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Thu, 18 Jun 2026 15:36:06 +0800 Subject: [PATCH 02/10] fix miad model --- ppmat/models/miad/__init__.py | 39 ----------- test/test_miad.py | 119 +++++++++++++++++++++++++++------- 2 files changed, 95 insertions(+), 63 deletions(-) diff --git a/ppmat/models/miad/__init__.py b/ppmat/models/miad/__init__.py index 6154349b..2a5bb807 100644 --- a/ppmat/models/miad/__init__.py +++ b/ppmat/models/miad/__init__.py @@ -12,47 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ppmat.models.miad.collate import CrystalBatch -from ppmat.models.miad.collate import MiADCollator -from ppmat.models.miad.collate import create_miad_dataloader -from ppmat.models.miad.collate import create_sampling_batch -from ppmat.models.miad.crystal_diffusion import CrystalGen as MiadCrystalGen -from ppmat.models.miad.crystal_diffusion import DiffCSP as MiadDiffCSP -from ppmat.models.miad.crystal_diffusion import init_diffusion -from ppmat.models.miad.crystal_diffusion import parse_batch -from ppmat.models.miad.diffusion_utils import TimeDistribution as MiadTimeDistribution -from ppmat.models.miad.frac_diffusion import PFM as MiadPFM -from ppmat.models.miad.frac_diffusion import WrappedNormal as MiadWrappedNormal -from ppmat.models.miad.lattice_diffusion import DDPM as MiadDDPM -from ppmat.models.miad.lattice_diffusion import FM as MiadFM -from ppmat.models.miad.lattice_diffusion import FM_LenAng as MiadFM_LenAng from ppmat.models.miad.miad import MiAD -from ppmat.models.miad.miad_cspnet import CSPNet as MiadCSPNet -from ppmat.models.miad.type_diffusion import D3PM as MiadD3PM -from ppmat.models.miad.type_diffusion import DDPM_onehot as MiadDDPM_onehot -from ppmat.schedulers.miad_schedulers import scheduler __all__ = [ - # Model "MiAD", - "MiadCSPNet", - "MiadCrystalGen", - "MiadDiffCSP", - "init_diffusion", - "parse_batch", - # Diffusion components - "MiadDDPM", - "MiadFM", - "MiadFM_LenAng", - "MiadWrappedNormal", - "MiadPFM", - "MiadDDPM_onehot", - "MiadD3PM", - "MiadTimeDistribution", - "scheduler", - # Collate - "MiADCollator", - "CrystalBatch", - "create_miad_dataloader", - "create_sampling_batch", ] diff --git a/test/test_miad.py b/test/test_miad.py index 23093ed5..2885be8f 100644 --- a/test/test_miad.py +++ b/test/test_miad.py @@ -20,10 +20,86 @@ from ppmat.datasets.mp20_dataset import MP20Dataset from ppmat.models import build_model from ppmat.models.miad.collate import MiADCollator # noqa: F401 +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": "diffcsp_cosine"}, + "frac_diffusion": {"method": "wrapped_normal", "scheduler": "default_wrapped_normal"}, + "type_diffusion": {"method": "d3pm", "scheduler": "default_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 that MiAD config can be parsed and model constructed.""" + """Test MiAD construction via build_model from yaml config.""" def test_config_load(self): config = OmegaConf.load("structure_generation/configs/miad/miad_mp20.yaml") @@ -33,42 +109,37 @@ def test_config_load(self): self.assertIn("diffusion_cfg", config["Model"]["__init_params__"]) self.assertIn("model_cfg", config["Model"]["__init_params__"]) - def test_model_construction(self): + 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.assertIsNotNone(model) + self.assertIsInstance(model, MiAD) self.assertIsInstance(model, paddle.nn.Layer) - def test_model_forward(self): + 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_size = 2 - num_atoms = paddle.to_tensor([5, 7], dtype="int64") - total_atoms = int(num_atoms.sum()) - batch_idx = paddle.concat( - [paddle.full([int(n)], i, dtype="int64") for i, n in enumerate(num_atoms)] - ) - 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") - - batch = { - "x0": [lattices, frac_coords, atom_types], - "batch_size": batch_size, - "num_atoms": num_atoms, - "batch_idx": batch_idx, - "atom_types": atom_types, - } - + 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): From e89e9d58a1b05a64e5a381e30d3b5a599c1c8a4a Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Mon, 22 Jun 2026 17:37:11 +0800 Subject: [PATCH 03/10] fix miad model --- ppmat/datasets/collate_fn.py | 1 - ppmat/metrics/prerelax_chgnet.py | 167 +----------- ppmat/models/miad/collate.py | 237 ------------------ ppmat/models/miad/crystal_diffusion.py | 123 +++------ ppmat/models/miad/default_configs.py | 88 ------- ppmat/models/miad/diffusion_utils.py | 9 +- ppmat/models/miad/frac_diffusion.py | 45 +--- ppmat/models/miad/graph_utils.py | 52 ---- ppmat/models/miad/lattice_diffusion.py | 51 +--- ppmat/models/miad/miad.py | 58 ++++- ppmat/models/miad/miad_cspnet.py | 133 ---------- ppmat/models/miad/type_diffusion.py | 28 +-- ppmat/schedulers/__init__.py | 2 - ppmat/schedulers/miad_schedulers.py | 1 - .../configs/miad/miad_mp20.yaml | 8 +- test/test_miad.py | 6 +- 16 files changed, 121 insertions(+), 888 deletions(-) delete mode 100644 ppmat/models/miad/collate.py delete mode 100644 ppmat/models/miad/default_configs.py delete mode 100644 ppmat/models/miad/graph_utils.py delete mode 100644 ppmat/models/miad/miad_cspnet.py diff --git a/ppmat/datasets/collate_fn.py b/ppmat/datasets/collate_fn.py index 69b68e19..9073af4c 100644 --- a/ppmat/datasets/collate_fn.py +++ b/ppmat/datasets/collate_fn.py @@ -30,7 +30,6 @@ from ppmat.datasets.custom_data_type import ConcatNumpyWarper from ppmat.datasets.geometric_data_type.batch import Batch from ppmat.datasets.geometric_data_type.data import Data -from ppmat.models.miad.collate import MiADCollator # noqa: F401 class DefaultCollator(object): diff --git a/ppmat/metrics/prerelax_chgnet.py b/ppmat/metrics/prerelax_chgnet.py index 2b090833..0e0bcb7a 100644 --- a/ppmat/metrics/prerelax_chgnet.py +++ b/ppmat/metrics/prerelax_chgnet.py @@ -13,18 +13,12 @@ # limitations under the License. """ -Pre-relaxation module using ppmat's built-in CHGNet for energy prediction. - -Provides structure relaxation and energy calculation needed for the -Stability metric in S.U.N. evaluation. Uses ppmat's CHGNet model -(Paddle implementation) -- no external PyTorch dependencies required. +CHGNet energy prediction utility for S.U.N. Stability metric. """ import logging -from typing import Dict from typing import List from typing import Optional -from typing import Tuple import numpy as np from pymatgen.core import Structure @@ -87,162 +81,3 @@ def predict_energy( energies.append(None) return energies - - -def relax_structure( - structure: Structure, - model_name: str = "chgnet_mptrj", - steps: int = 200, - fmax: float = 0.05, - step_size: float = 0.02, -) -> Tuple[Optional[Structure], Optional[float], int]: - """Relax a crystal structure using gradient-based optimization with CHGNet. - - Performs iterative structure relaxation by computing forces from CHGNet - energy predictions and updating atomic positions accordingly. - - Args: - structure: pymatgen Structure to relax. - model_name: CHGNet model name in ppmat registry. - steps: Maximum number of relaxation steps. - fmax: Force convergence threshold (eV/A). - step_size: Position update step size. - - Returns: - Tuple of (relaxed_structure, energy_per_atom, n_steps). - None values if relaxation fails. - """ - model, converter = _load_chgnet(model_name) - - try: - import paddle - - current_struct = structure.copy() - converged_step = steps - - for step in range(steps): - graph = converter(current_struct) - pred = model.predict(graph) - e_per_atom = float(np.array(pred["energy_per_atom"]).flatten()[0]) - - forces = np.array(pred["force"]) - max_force = np.max(np.abs(forces)) - - if max_force < fmax: - converged_step = step + 1 - break - - # Simple gradient descent on positions - frac_coords = current_struct.frac_coords - lattice_matrix = current_struct.lattice.matrix - - # Force in Cartesian -> displacement in fractional - cart_displacements = -step_size * forces / (max_force + 1e-8) - frac_displacements = cart_displacements @ np.linalg.inv(lattice_matrix) - - new_frac = frac_coords + frac_displacements - # Wrap to [0, 1) - new_frac = new_frac % 1.0 - - current_struct = Structure( - current_struct.lattice, - current_struct.species, - new_frac, - coords_are_cartesian=False, - ) - - # Final energy - graph = converter(current_struct) - pred = model.predict(graph) - final_energy = float(np.array(pred["energy_per_atom"]).flatten()[0]) - - return current_struct, final_energy, converged_step - - except Exception as e: - logger.debug(f"Relaxation failed: {e}") - return None, None, 0 - - -def predict_energies_with_relaxation( - structures: List[Optional[Structure]], - model_name: str = "chgnet_mptrj", - steps: int = 200, - fmax: float = 0.05, - relax: bool = True, -) -> Dict[str, List]: - """Predict energies for structures, optionally with pre-relaxation. - - This is the main entry point for S.U.N. Stability computation. - - Args: - structures: List of pymatgen Structures (None for invalid). - model_name: CHGNet model name in ppmat registry. - steps: Maximum relaxation steps. - fmax: Force convergence threshold (eV/A). - relax: Whether to relax structures before energy evaluation. - - Returns: - Dict with: - - 'energy_gen': energy per atom before relaxation - - 'energy_relaxed': energy per atom after relaxation (or same as gen) - - 'relaxed_structures': relaxed Structures (or originals) - - 'num_sites': number of atoms per structure - """ - energy_gen = [] - energy_relaxed = [] - relaxed_structures = [] - num_sites = [] - - for struct in tqdm( - structures, - desc=f"{'Relaxing' if relax else 'Evaluating'} structures (CHGNet)", - ): - if struct is None: - energy_gen.append(None) - energy_relaxed.append(None) - relaxed_structures.append(None) - num_sites.append(None) - continue - - try: - n_sites = struct.num_sites - model, converter = _load_chgnet(model_name) - - # Initial energy - graph = converter(struct) - pred = model.predict(graph) - e_gen = float(np.array(pred["energy_per_atom"]).flatten()[0]) * n_sites - - if relax: - relaxed, e_relax, n_steps = relax_structure( - struct, model_name=model_name, steps=steps, fmax=fmax - ) - if relaxed is not None: - energy_gen.append(e_gen / n_sites) - energy_relaxed.append(e_relax) - relaxed_structures.append(relaxed) - num_sites.append(n_sites) - else: - energy_gen.append(e_gen / n_sites) - energy_relaxed.append(e_gen / n_sites) - relaxed_structures.append(struct) - num_sites.append(n_sites) - else: - energy_gen.append(e_gen / n_sites) - energy_relaxed.append(e_gen / n_sites) - relaxed_structures.append(struct) - num_sites.append(n_sites) - - except Exception as e: - logger.debug(f"Structure evaluation failed: {e}") - energy_gen.append(None) - energy_relaxed.append(None) - relaxed_structures.append(None) - num_sites.append(None) - - return { - "energy_gen": energy_gen, - "energy_relaxed": energy_relaxed, - "relaxed_structures": relaxed_structures, - "num_sites": num_sites, - } diff --git a/ppmat/models/miad/collate.py b/ppmat/models/miad/collate.py deleted file mode 100644 index 5b7ba9bc..00000000 --- a/ppmat/models/miad/collate.py +++ /dev/null @@ -1,237 +0,0 @@ -# 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 typing import Any -from typing import Dict -from typing import List -from typing import Optional - -import numpy as np -import paddle - -from ppmat.datasets.custom_data_type import ConcatData - - -def _extract_concat_data(data): - if data is None: - return None - if isinstance(data, ConcatData): - return np.array(data.data) - return data - - -class CrystalBatch: - """Minimal batch object mimicking torch_geometric Batch interface.""" - - def __init__( - self, - num_atoms: paddle.Tensor, - atom_types: paddle.Tensor, - batch: paddle.Tensor, - frac_coords: Optional[paddle.Tensor] = None, - lattice: Optional[paddle.Tensor] = None, - ): - self.num_atoms = num_atoms - self.atom_types = atom_types - self.batch = batch - self.frac_coords = frac_coords - self.lattice = lattice - - def cuda(self, device: str = "gpu"): - place = paddle.CUDAPlace(0) if device == "gpu" else paddle.CPUPlace() - self.num_atoms = self.num_atoms.to(place) - self.atom_types = self.atom_types.to(place) - self.batch = self.batch.to(place) - if self.frac_coords is not None: - self.frac_coords = self.frac_coords.to(place) - if self.lattice is not None: - self.lattice = self.lattice.to(place) - return self - - -class MiADCollator: - """Collate function for MiAD model. - - Converts MP20Dataset batch to CrystalGen format. - """ - - def __init__( - self, - mirage_max_atoms: Optional[int] = None, - use_mirage: bool = False, - ): - self.mirage_max_atoms = mirage_max_atoms - self.use_mirage = use_mirage - self.N_m = mirage_max_atoms - - def __call__(self, batch_list: List[Dict]) -> Dict[str, Any]: - if not batch_list: - return {} - batch_size = len(batch_list) - all_frac_coords = [] - all_atom_types = [] - all_num_atoms = [] - all_lattices = [] - batch_idx_list = [] - - for i, sample_i in enumerate(batch_list): - if not isinstance(sample_i, dict): - continue - if "structure_array" in sample_i: - sa = sample_i["structure_array"] - elif "frac_coords" in sample_i: - sa = sample_i - else: - continue - - frac_coords = _extract_concat_data(sa.get("frac_coords")) - atom_types = _extract_concat_data(sa.get("atom_types")) - lattice = _extract_concat_data(sa.get("lattice")) - num_atoms = _extract_concat_data(sa.get("num_atoms")) - - if frac_coords is None or atom_types is None or lattice is None: - continue - - if frac_coords.ndim >= 2: - n_atoms = frac_coords.shape[0] - elif num_atoms is not None: - if isinstance(num_atoms, np.ndarray) and num_atoms.ndim > 0: - n_atoms = int(num_atoms.flatten()[0]) - else: - n_atoms = int(num_atoms) - else: - continue - - if frac_coords.ndim == 1: - frac_coords = frac_coords.reshape(-1, 3) - if atom_types.ndim > 1: - atom_types = atom_types.flatten() - - if lattice.ndim == 2: - lattice = lattice.reshape(1, 3, 3) - elif lattice.ndim == 1: - lattice = lattice.reshape(1, 3, 3) - - if self.N_m is not None and self.use_mirage and n_atoms < self.N_m: - frac_coords, atom_types = self._pad_mirage( - frac_coords, atom_types, n_atoms, self.N_m - ) - n_atoms = self.N_m - - all_frac_coords.append(frac_coords) - all_atom_types.append(atom_types) - all_lattices.append(lattice[0] if lattice.shape[0] == 1 else lattice) - all_num_atoms.append(n_atoms) - batch_idx_list.extend([i] * n_atoms) - - if not all_frac_coords: - return {} - - frac_coords_cat = np.concatenate(all_frac_coords, axis=0).astype("float32") - atom_types_cat = np.concatenate(all_atom_types, axis=0).astype("int64") - lattices_cat = np.stack(all_lattices, axis=0).astype("float32") - num_atoms_arr = np.array(all_num_atoms, dtype="int64") - batch_idx = np.array(batch_idx_list, dtype="int64") - - frac_coords_t = paddle.to_tensor(frac_coords_cat) - atom_types_t = paddle.to_tensor(atom_types_cat) - lattices_t = paddle.to_tensor(lattices_cat) - num_atoms_t = paddle.to_tensor(num_atoms_arr) - batch_idx_t = paddle.to_tensor(batch_idx) - - crystal_batch = CrystalBatch( - num_atoms=num_atoms_t, - atom_types=atom_types_t, - batch=batch_idx_t, - frac_coords=frac_coords_t, - lattice=lattices_t, - ) - - return { - "x0": [lattices_t, frac_coords_t, atom_types_t], - "batch_size": batch_size, - "num_atoms": num_atoms_t, - "batch_idx": batch_idx_t, - "atom_types": atom_types_t, - "batch": crystal_batch, - } - - @staticmethod - def _pad_mirage(frac_coords, atom_types, n_atoms, N_m): - pad_n = N_m - n_atoms - if pad_n > 0: - pad_fc = np.random.rand(pad_n, 3).astype("float32") - pad_at = np.zeros(pad_n, dtype="int64") - frac_coords = np.concatenate([frac_coords, pad_fc], axis=0) - atom_types = np.concatenate([atom_types, pad_at], axis=0) - return frac_coords, atom_types - - -def create_miad_dataloader( - dataset, - batch_size: int, - shuffle: bool = False, - num_workers: int = 0, - mirage_max_atoms: Optional[int] = None, - use_mirage: bool = False, - **kwargs, -) -> paddle.io.DataLoader: - collate_fn = MiADCollator( - mirage_max_atoms=mirage_max_atoms, - use_mirage=use_mirage, - ) - return paddle.io.DataLoader( - dataset, - batch_size=batch_size, - shuffle=shuffle, - num_workers=num_workers, - collate_fn=collate_fn, - **kwargs, - ) - - -def create_sampling_batch( - batch_size: int, - num_atoms: Optional[paddle.Tensor] = None, - mirage_max_atoms: Optional[int] = None, - device: str = "cpu", -) -> Dict[str, Any]: - if mirage_max_atoms is not None: - N_m = mirage_max_atoms - num_atoms = paddle.full([batch_size], N_m, dtype="int64") - total_atoms = batch_size * N_m - elif num_atoms is not None: - total_atoms = int(num_atoms.sum()) - else: - num_atoms = paddle.randint(5, 30, [batch_size], dtype="int64") - total_atoms = int(num_atoms.sum()) - - node2graph = paddle.repeat_interleave( - paddle.arange(batch_size, dtype="int64"), num_atoms - ) - - crystal_batch = CrystalBatch( - num_atoms=num_atoms, - atom_types=paddle.zeros([total_atoms], dtype="int64"), - batch=node2graph, - frac_coords=paddle.zeros([total_atoms, 3], dtype="float32"), - lattice=paddle.zeros([batch_size, 3, 3], dtype="float32"), - ) - - return { - "x0": [None, None, None], - "batch_size": batch_size, - "num_atoms": num_atoms, - "batch": crystal_batch, - } diff --git a/ppmat/models/miad/crystal_diffusion.py b/ppmat/models/miad/crystal_diffusion.py index c29a2108..4ab55a98 100644 --- a/ppmat/models/miad/crystal_diffusion.py +++ b/ppmat/models/miad/crystal_diffusion.py @@ -45,31 +45,13 @@ def parse_num_atoms_to_per_crystal(num_atoms_data): return paddle.to_tensor(num_atoms_np.astype("int64")), num_atoms_np -def parse_batch(batch): - if "batch" in batch: - return { - "num_atoms": batch["batch"].num_atoms, - "batch_idx": batch["batch"].batch, - "atom_types": batch["batch"].atom_types, - "batch_size": batch["batch_size"], - } - if "batch_idx" in batch: - return { - "num_atoms": batch["num_atoms"], - "batch_idx": batch["batch_idx"], - "atom_types": batch["atom_types"], - "batch_size": batch.get("batch_size", len(batch["num_atoms"])), - } - raise ValueError("Cannot determine batch format: missing 'batch' or 'batch_idx'") def init_diffusion(diffusion_config, logger): - if hasattr(diffusion_config, "default_config") and diffusion_config.default_config: - diffusion_config = _apply_default_config(diffusion_config.default_config) switch = { "Default": CrystalGen, - "DiffCSP": DiffCSP, + "DiffCSP": CrystalGen, } method = diffusion_config.method if method in switch: @@ -77,14 +59,6 @@ def init_diffusion(diffusion_config, logger): raise NotImplementedError(f"Diffusion method '{method}' not implemented") -def _apply_default_config(config_name): - from ppmat.models.miad.default_configs import DEFAULT_DIFFUSION_CONFIGS - - if config_name in DEFAULT_DIFFUSION_CONFIGS: - return DEFAULT_DIFFUSION_CONFIGS[config_name] - raise KeyError(f"Unknown default config: '{config_name}'") - - class CrystalGen: """Standard crystal generation with single-step reverse sampling.""" @@ -140,10 +114,18 @@ def forward_step_sample(self, x0, t, batch): def reverse_step_sample(self, xt, t, model, batch): lt, ft, at = xt - batch["prediction"] = self.model_prediction(xt, t, model, batch) - l_pred, f_pred, a_pred = batch["prediction"] - lt_1 = self.lat_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) - ft_1 = self.frac_diffusion.reverse_step_sample(f_pred, ft, t[1], batch) + if self.config.method == "DiffCSP": + _, f_pred, _ = self.model_prediction(xt, t, model, batch) + ft_05 = self.frac_diffusion.reverse_step_sample_part_1(f_pred, ft, t[1], batch) + xt_05 = [lt, ft_05, at] + l_pred, f_pred, a_pred = self.model_prediction(xt_05, t, model, batch) + lt_1 = self.lat_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) + ft_1 = self.frac_diffusion.reverse_step_sample_part_2(f_pred, ft_05, t[1], batch) + else: + batch["prediction"] = self.model_prediction(xt, t, model, batch) + l_pred, f_pred, a_pred = batch["prediction"] + lt_1 = self.lat_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) + ft_1 = self.frac_diffusion.reverse_step_sample(f_pred, ft, t[1], batch) at_1 = ( self.type_diffusion.reverse_step_sample(a_pred, at, t[1], batch) if self.gen_type @@ -152,10 +134,21 @@ def reverse_step_sample(self, xt, t, model, batch): return [lt_1, ft_1, at_1] def _get_batch_info(self, batch): - try: - return parse_batch(batch) - except ValueError: - return self._normalize_batch(batch) + if "batch" in batch: + return { + "num_atoms": batch["batch"].num_atoms, + "batch_idx": batch["batch"].batch, + "atom_types": batch["batch"].atom_types, + "batch_size": batch["batch_size"], + } + if "batch_idx" in batch: + return { + "num_atoms": batch["num_atoms"], + "batch_idx": batch["batch_idx"], + "atom_types": batch["atom_types"], + "batch_size": batch.get("batch_size", len(batch["num_atoms"])), + } + return self._normalize_batch(batch) def _normalize_batch(self, batch): if "structure_array" in batch: @@ -197,35 +190,13 @@ def prior_sample(self, batch): def model_prediction(self, xt, t, model, batch): batch_info = self._get_batch_info(batch) lt, ft, at = xt - nn_pred = model( - t[0], - self.time_embedding(1000 * (t[0] / self.num_steps) + 1), - at, - ft, - lt, - batch_info["num_atoms"], - batch_info["batch_idx"], - ) + time_emb = self.time_embedding(1000 * (t[0] / self.num_steps) + 1) + nn_pred = model(time_emb, at, ft, lt, batch_info["num_atoms"], batch_info["batch_idx"]) l_pred = nn_pred[0] f_pred = nn_pred[1] a_pred = nn_pred[2] if self.gen_type else None return [l_pred, f_pred, a_pred] - def get_x0_prediction(self, pred, xt, t, batch): - l_pred, f_pred, a_pred = pred - lt, ft, at = xt - return [ - self.lat_diffusion.get_x0_prediction(l_pred, lt, t[0], batch), - self.frac_diffusion.get_x0_prediction(f_pred, ft, t[1], batch), - ( - self.type_diffusion.get_x0_prediction( - a_pred, at, t[1], batch, x0_format="disc" - ) - if self.gen_type - else at.clone().detach() - ), - ] - def train_step(self, batch, model, mode): modifications = os.environ.get("MODIFICATIONS_FIELD", "") @@ -309,7 +280,7 @@ def sampling_procedure(self, model, batch, progress_printer): batch, start_from=self.num_steps - 1 ) for t_vector in reverse_time_iterator: - batch["t"] = self.time_distribution.to_cuda(t_vector, batch) + batch["t"] = t_vector t_value = batch["t"][0][0].item() progress_printer(t_value) batch["xt"] = self.reverse_step_sample( @@ -320,37 +291,13 @@ def sampling_procedure(self, model, batch, progress_printer): 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 [ - self.lat_diffusion.output_transform(x0[0], batch), - self.frac_diffusion.output_transform(x0[1], batch), + _out(self.lat_diffusion, x0[0]), + _out(self.frac_diffusion, x0[1]), ( - self.type_diffusion.output_transform(x0[2], batch) + _out(self.type_diffusion, x0[2]) if self.gen_type else x0[2] ), ] - - -class DiffCSP(CrystalGen): - """DiffCSP with two-step reverse sampling.""" - - def reverse_step_sample(self, xt, t, model, batch): - lt, ft, at = xt - - # Step 1: only use fractional prediction - _, f_pred, _ = self.model_prediction(xt, t, model, batch) - ft_05 = self.frac_diffusion.reverse_step_sample_part_1(f_pred, ft, t[1], batch) - xt_05 = [lt, ft_05, at] - - # Step 2: full prediction with corrected fractional coords - l_pred, f_pred, a_pred = self.model_prediction(xt_05, t, model, batch) - lt_1 = self.lat_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) - ft_1 = self.frac_diffusion.reverse_step_sample_part_2( - f_pred, ft_05, t[1], batch - ) - 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] diff --git a/ppmat/models/miad/default_configs.py b/ppmat/models/miad/default_configs.py deleted file mode 100644 index c400ba8f..00000000 --- a/ppmat/models/miad/default_configs.py +++ /dev/null @@ -1,88 +0,0 @@ -# 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. - - -class _Cfg: - def __init__(self, **kwargs): - for k, v in kwargs.items(): - if isinstance(v, dict): - setattr(self, k, _Cfg(**v)) - else: - setattr(self, k, v) - - -def _make_config(method, lat_diffusion, frac_diffusion, type_diffusion): - return _Cfg( - method=method, - cont_time=False, - num_steps=1000, - task="", - lat_diffusion=lat_diffusion, - frac_diffusion=frac_diffusion, - type_diffusion=type_diffusion, - ) - - -_lat = { - "diffcsp_ddpm_config": _Cfg(method="ddpm", scheduler="diffcsp_cosine"), - "ddpm_config": _Cfg(method="ddpm", scheduler="cosine"), - "fm_config": _Cfg(method="fm", parameterization="eps"), - "fm_v_config": _Cfg(method="fm", parameterization="v"), - "fm_lenang_config": _Cfg(method="fm_lenang"), -} - -_coord = { - "wrapped_normal_config": _Cfg( - method="wrapped_normal", scheduler="default_wrapped_normal" - ), - "pfm_config": _Cfg(method="pfm"), -} - -_type = { - "d3pm_config": _Cfg(method="d3pm", scheduler="default_d3pm"), - "ddpm_onehot_config": _Cfg(method="ddpm_onehot", scheduler="diffcsp_cosine"), -} - -DEFAULT_DIFFUSION_CONFIGS = { - "DiffCSP:diffcsp_ddpm_config:wrapped_normal_config:d3pm_config": _make_config( - "DiffCSP", - _lat["diffcsp_ddpm_config"], - _coord["wrapped_normal_config"], - _type["d3pm_config"], - ), - "Default:ddpm_config:wrapped_normal_config:d3pm_config": _make_config( - "Default", - _lat["ddpm_config"], - _coord["wrapped_normal_config"], - _type["d3pm_config"], - ), - "Default:fm_config:wrapped_normal_config:d3pm_config": _make_config( - "Default", - _lat["fm_config"], - _coord["wrapped_normal_config"], - _type["d3pm_config"], - ), - "DiffCSP:ddpm_config:wrapped_normal_config:ddpm_onehot_config": _make_config( - "DiffCSP", - _lat["ddpm_config"], - _coord["wrapped_normal_config"], - _type["ddpm_onehot_config"], - ), - "DiffCSP:fm_config:wrapped_normal_config:d3pm_config": _make_config( - "DiffCSP", - _lat["fm_config"], - _coord["wrapped_normal_config"], - _type["d3pm_config"], - ), -} diff --git a/ppmat/models/miad/diffusion_utils.py b/ppmat/models/miad/diffusion_utils.py index db122066..3ceb6acc 100644 --- a/ppmat/models/miad/diffusion_utils.py +++ b/ppmat/models/miad/diffusion_utils.py @@ -76,10 +76,6 @@ def reverse_time_iterator(self, batch, start_from=-1): t = paddle.full([batch["batch_size"]], t_val, dtype="float32") yield self._atom_expand(batch, t) - def to_cuda(self, t_vector, batch): - return t_vector - - def mean_interleave(t, num_repeats): """Compute per-group mean using paddle_scatter.scatter_mean. @@ -94,3 +90,8 @@ def mean_interleave(t, num_repeats): paddle.arange(num_repeats.shape[0]), num_repeats ) return scatter_mean(t, batch_idx, dim=0) + + +def total_atoms_from_batch(batch): + num_atoms = batch["num_atoms"] + return int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) diff --git a/ppmat/models/miad/frac_diffusion.py b/ppmat/models/miad/frac_diffusion.py index 392b4f22..8243028d 100644 --- a/ppmat/models/miad/frac_diffusion.py +++ b/ppmat/models/miad/frac_diffusion.py @@ -21,6 +21,7 @@ import paddle from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler +from ppmat.models.miad.diffusion_utils import total_atoms_from_batch class WrappedNormal: @@ -55,13 +56,8 @@ def __init__(self, diffusion_config): if self.step_lr is None: self.step_lr = _DEFAULT_GAMMA.get(diffusion_config.task, 1e-5) - self.drift_step_coef = paddle.ones([self.num_steps], dtype="float32")[:, None] - self.diff_step_coef = paddle.ones([self.num_steps], dtype="float32")[:, None] self.to_domain = lambda x: x - def output_transform(self, x0, batch): - return x0 - def forward_step_sample(self, x0, t, batch): st = self.sigmas_t[t.cast("int64")] self.randn_x = paddle.randn(x0.shape) @@ -74,13 +70,8 @@ def reverse_step_sample_part_1(self, normed_score_pred, xt, t, batch): snt = self.sigmas_norm_t[t_idx] step_size = self.step_lr * (st / self.sb) ** 2 std_x = paddle.sqrt(2 * step_size) - drift = ( - -step_size - * normed_score_pred - * paddle.sqrt(snt) - * self.drift_step_coef[t_idx] - ) - diffusion = std_x * paddle.randn(xt.shape) * self.diff_step_coef[t_idx] + drift = -step_size * normed_score_pred * paddle.sqrt(snt) + diffusion = std_x * paddle.randn(xt.shape) xt_05 = xt + drift + diffusion return xt_05 @@ -91,13 +82,8 @@ def reverse_step_sample_part_2(self, normed_score_pred, xt_05, t, batch): snt = self.sigmas_norm_t[t_idx] step_size = st**2 - st_1**2 std_x = paddle.sqrt((st_1**2 * (st**2 - st_1**2)) / (st**2)) - drift = ( - -step_size - * normed_score_pred - * paddle.sqrt(snt) - * self.drift_step_coef[t_idx] - ) - diffusion = std_x * paddle.randn(xt_05.shape) * self.diff_step_coef[t_idx] + drift = -step_size * normed_score_pred * paddle.sqrt(snt) + diffusion = std_x * paddle.randn(xt_05.shape) xt_1 = (xt_05 + drift + diffusion) % 1.0 return xt_1 @@ -107,8 +93,7 @@ def reverse_step_sample(self, normed_score_pred, xt, t, batch): return xt_1 def prior_sample(self, batch): - num_atoms = batch["num_atoms"] - total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + total_atoms = total_atoms_from_batch(batch) return paddle.rand([total_atoms, 3], dtype="float32") def loss(self, batch): @@ -134,10 +119,6 @@ def loss(self, batch): return l2 - def get_x0_prediction(self, normed_score_pred, xt, t, batch): - t_idx = t.cast("int64") - x0_pred = (xt - self.sigmas_t[t_idx] * normed_score_pred) % 1.0 - return x0_pred class PFM: @@ -148,14 +129,10 @@ def __init__(self, diffusion_config): self.num_steps = diffusion_config.num_steps self.cont_time = diffusion_config.cont_time self.final_t = diffusion_config.num_steps - self.step_coef = paddle.ones([self.final_t], dtype="float32")[:, None] self.step = paddle.to_tensor(1.0 / self.final_t) self.to_domain = lambda x: x self.default_loss_scale = 10 - def output_transform(self, x0, batch): - return x0 - def forward_step_sample(self, x0, t, batch): xT_minus_x0 = paddle.rand(x0.shape) - 0.5 step = (1 + t[:, None]) * self.step @@ -164,20 +141,14 @@ def forward_step_sample(self, x0, t, batch): return xt def reverse_step_sample(self, vt, xt, t, batch): - t_idx = t.cast("int64") - xt_1 = (xt + self.step_coef[t_idx] * self.step * vt) % 1.0 + xt_1 = (xt + self.step * vt) % 1.0 return xt_1 def prior_sample(self, batch): - num_atoms = batch["num_atoms"] - total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + total_atoms = total_atoms_from_batch(batch) return paddle.rand([total_atoms, 3]) def loss(self, batch): vt = batch["prediction"][1] l2 = ((vt - self.ut) ** 2).reshape([-1, 3]).mean(axis=1) return l2 * self.default_loss_scale - - def get_x0_prediction(self, vt, xt, t, batch): - step = (1 + t[:, None]) * self.step - return (xt + step * vt) % 1.0 diff --git a/ppmat/models/miad/graph_utils.py b/ppmat/models/miad/graph_utils.py deleted file mode 100644 index 150618e9..00000000 --- a/ppmat/models/miad/graph_utils.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. - -"""Graph utility for converting dense adjacency to COO edge_index.""" - -import paddle - - -def dense_to_sparse(adj): - """Convert a dense adjacency matrix to sparse (edge_index, edge_attr) COO format. - - Supports both single-graph (2D) and batched (3D) dense adjacency tensors. - """ - if adj.ndim == 2: - nonzero = paddle.nonzero(adj.cast("float32")) - if nonzero.shape[0] == 0: - return paddle.zeros([2, 0], dtype="int64"), paddle.zeros( - [0], dtype="float32" - ) - edge_index = nonzero.t() - return edge_index, paddle.ones([edge_index.shape[1]], dtype="float32") - - # Batched case - batch_size, num_nodes, _ = adj.shape - edge_indices_list = [] - edge_attrs_list = [] - - for b in range(batch_size): - adj_b = adj[b] - nonzero = paddle.nonzero(adj_b.cast("float32")) - if nonzero.shape[0] == 0: - continue - edge_indices_list.append(nonzero.t()) - edge_attrs_list.append(paddle.ones([nonzero.shape[0]], dtype="float32")) - - if not edge_indices_list: - return paddle.zeros([2, 0], dtype="int64"), paddle.zeros([0], dtype="float32") - - edge_index = paddle.concat(edge_indices_list, axis=1) - edge_attr = paddle.concat(edge_attrs_list, axis=0) - return edge_index, edge_attr diff --git a/ppmat/models/miad/lattice_diffusion.py b/ppmat/models/miad/lattice_diffusion.py index f0af3670..02e14f8a 100644 --- a/ppmat/models/miad/lattice_diffusion.py +++ b/ppmat/models/miad/lattice_diffusion.py @@ -67,15 +67,8 @@ def __init__(self, diffusion_config, scheduler_cfg=None): self.betas_t * (1 - self.cumprod_alphas_t_1) / (1 - self.cumprod_alphas_t) ) - # Pre-compute eps-to-x0 coefficients - self.eps_to_x0_c0 = paddle.sqrt(1 / self.cumprod_alphas_t) - self.eps_to_x0_c1 = paddle.sqrt(1 / self.cumprod_alphas_t - 1) - self.to_domain = lambda x: x - def output_transform(self, x0, batch): - return x0 - def forward_step_sample(self, x0, t, batch): if self.cont_time and self.f_cumprod_alphas_t is not None: at = self.f_cumprod_alphas_t(t) @@ -104,11 +97,6 @@ def loss(self, batch): ) return l2 - def get_x0_prediction(self, eps_pred, xt, t, batch): - t_idx = t.cast("int64") - x0_pred = self.eps_to_x0_c0[t_idx] * xt - self.eps_to_x0_c1[t_idx] * eps_pred - return x0_pred - class FM: """Flow Matching for lattice diffusion.""" @@ -121,12 +109,8 @@ def __init__(self, diffusion_config): self.parameterization = getattr( diffusion_config.lat_diffusion, "parameterization", "eps" ) - self.step_coef = paddle.ones([self.num_steps], dtype="float32")[:, None, None] self.to_domain = lambda x: x - def output_transform(self, x0, batch): - return x0 - def forward_step_sample(self, x0, t, batch): eps = paddle.randn(x0.shape) step = (1 + t[:, None, None]) * self.step @@ -148,7 +132,7 @@ def reverse_step_sample(self, pred, xt, t, batch): vt = x0_pred - eps_pred elif self.parameterization == "v": vt = pred - xt_1 = xt + self.step_coef[t_idx] * self.step * vt + xt_1 = xt + self.step * vt return xt_1 def prior_sample(self, batch): @@ -159,19 +143,6 @@ def loss(self, batch): l2 = ((vt - self.ut) ** 2).reshape([-1, 9]).mean(axis=1) return l2 - def get_x0_prediction(self, pred, xt, t, batch): - step = (1 + t[:, None, None]) * self.step - t_idx = t.cast("int64") - if self.parameterization == "eps": - if t_idx[0] >= 999: - return paddle.randn(xt.shape) - eps_pred = (-1) * pred - x0_pred = (xt - step * eps_pred) / (1 - step) - vt = x0_pred - eps_pred - elif self.parameterization == "v": - vt = pred - return xt + step * vt - class FM_LenAng: """Flow Matching for lattice in length-angle parameterization.""" @@ -188,25 +159,17 @@ def __init__(self, diffusion_config): self.angle_difference_bound = 20 self.default_loss_scale = 0.45 - self._lat2lenang = None self._lenang2lat = None self.to_domain = lambda x: x - def _get_converters(self): - if self._lat2lenang is None: + def _get_converter(self): + if self._lenang2lat is None: from ppmat.utils.crystal import lattice_params_to_matrix_paddle - from ppmat.utils.crystal import lattices_to_params_shape_paddle - self._lat2lenang = lambda lat: paddle.concat( - lattices_to_params_shape_paddle(lat), axis=1 - ) self._lenang2lat = lambda la: lattice_params_to_matrix_paddle( la[:, :3], la[:, 3:] ) - return self._lat2lenang, self._lenang2lat - - def output_transform(self, x0, batch): - return x0 + return self._lenang2lat def forward_step_sample(self, x0, t, batch): xT = self.prior_sample(batch) @@ -239,7 +202,7 @@ def prior_sample(self, batch): valid = ang[check][:bs] lenang_xT[:, 3:] = valid - _, lenang2lat = self._get_converters() + lenang2lat = self._get_converter() xT = lenang2lat(lenang_xT) return xT @@ -247,7 +210,3 @@ def loss(self, batch): vt = batch["prediction"][0] l2 = ((vt - self.ut) ** 2).reshape([-1, 9]).mean(axis=1) return self.default_loss_scale * l2 - - def get_x0_prediction(self, vt, xt, t, batch): - step = (1 + t[:, None, None]) * self.step - return xt + step * vt diff --git a/ppmat/models/miad/miad.py b/ppmat/models/miad/miad.py index 9f21c94d..c04781c7 100644 --- a/ppmat/models/miad/miad.py +++ b/ppmat/models/miad/miad.py @@ -20,11 +20,44 @@ from ppmat.models.miad.crystal_diffusion import init_diffusion from ppmat.models.miad.crystal_diffusion import parse_num_atoms_to_per_crystal -from ppmat.models.miad.miad_cspnet import CSPNet +from ppmat.models.diffcsp.diffcsp import CSPNet from ppmat.utils import logger from ppmat.utils.crystal import lattices_to_params_shape_numpy +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 + + def _dict_to_sns(d): """Recursively convert dict to SimpleNamespace for attribute access.""" if not isinstance(d, dict): @@ -41,7 +74,27 @@ def __init__(self, model_cfg=None, diffusion_cfg=None, **kwargs): model_cfg = model_cfg or {} diffusion_cfg = diffusion_cfg or {} - self.decoder = CSPNet(**model_cfg) + 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 = CSPNet(**cspnet_kwargs) + # Remove unused prop_mlp (parent creates it; checkpoint lacks these weights) + for i in range(self.decoder.num_layers): + layer = getattr(self.decoder, f"csp_layer_{i}", None) + if layer and hasattr(layer, "prop_mlp"): + del layer.prop_mlp + # Replace gen_edges with block_diag (parent meshgrid causes GPU crash on certain batch sizes) + def _gen_edges(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]]) + self.decoder.gen_edges = _gen_edges if isinstance(diffusion_cfg, dict): diffusion_cfg = _dict_to_sns(diffusion_cfg) self.diffusion = init_diffusion(diffusion_cfg, logger=None) @@ -142,6 +195,7 @@ def set_state_dict(self, state_dict, use_structured_name=True): 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, diff --git a/ppmat/models/miad/miad_cspnet.py b/ppmat/models/miad/miad_cspnet.py deleted file mode 100644 index d5b9a278..00000000 --- a/ppmat/models/miad/miad_cspnet.py +++ /dev/null @@ -1,133 +0,0 @@ -# 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. - -""" -CSPNet for MiAD model. Inherits core structure from ppmat.models.diffcsp.diffcsp.CSPNet. -Overrides gen_edges for block_diag edge construction and forward for pre-computed -time embeddings and flexible atom type (one-hot / discrete) handling. -""" - -import paddle - -from paddle_scatter import scatter -from ppmat.models.diffcsp.diffcsp import CSPNet as _DiffCSPCSPNet -from ppmat.models.miad.graph_utils import dense_to_sparse - - -class CSPNet(_DiffCSPCSPNet): - """MiAD-specific CSPNet inheriting from DiffCSP CSPNet. - - Key differences from the parent: - - gen_edges: uses paddle.block_diag + dense_to_sparse for fully-connected edges - - forward: accepts pre-computed t_emb (time embedding) and handles both - one-hot and discrete atom types - - Removes prop_mlp from each CSPLayer (MiAD does not use property embeddings) - """ - - def __init__( - self, - hidden_dim=128, - latent_dim=256, - num_layers=4, - max_atoms=100, - act_fn="silu", - dis_emb="sin", - num_freqs=10, - edge_style="fc", - ln=False, - ip=True, - smooth=False, - pred_type=False, - cutoff=7.0, - max_neighbors=20, - model_name=None, - **kwargs, - ): - self.cutoff = cutoff - self.max_neighbors = max_neighbors - super().__init__( - hidden_dim=hidden_dim, - latent_dim=latent_dim, - num_layers=num_layers, - act_fn=act_fn, - dis_emb=dis_emb, - num_freqs=num_freqs, - edge_style=edge_style, - ln=ln, - ip=ip, - smooth=smooth, - pred_type=pred_type, - num_classes=max_atoms, - ) - for i in range(self.num_layers): - layer = self._modules.get(f"csp_layer_{i}") - if layer and hasattr(layer, 'prop_mlp'): - del layer.prop_mlp - - def gen_edges(self, num_atoms, frac_coords): - """Fully-connected graph via block diagonal adjacency matrix.""" - if self.edge_style == "fc": - lis = [paddle.ones([n, n], dtype=num_atoms.dtype) for n in num_atoms] - fc_graph = paddle.block_diag(lis) - fc_edges, _ = dense_to_sparse(fc_graph) - return fc_edges, (frac_coords[fc_edges[1]] - frac_coords[fc_edges[0]]) - - def forward( - self, t, t_emb, atom_types, frac_coords, lattices, num_atoms, node2graph - ): - edges, frac_diff = self.gen_edges(num_atoms, frac_coords) - edge2graph = node2graph[edges[0]] - # Handle both discrete atom types and one-hot encoding - if atom_types.ndim > 1: - if self.smooth: - node_features = self.node_embedding(atom_types.cast("float32")) - else: - atom_indices = atom_types.argmax(axis=-1) - node_features = self.node_embedding(atom_indices.cast("int64")) - else: - if self.smooth: - node_features = self.node_embedding(atom_types.cast("float32")) - else: - node_features = self.node_embedding(atom_types - 1) - - t_per_atom = paddle.repeat_interleave(t_emb, num_atoms, axis=0) - node_features = paddle.concat([node_features, t_per_atom], axis=1) - node_features = self.atom_latent_emb(node_features) - - for i in range(0, self.num_layers): - node_features = self._modules["csp_layer_%d" % i]( - node_features, - frac_coords, - lattices, - edges, - edge2graph, - frac_diff=frac_diff, - ) - - if self.ln: - node_features = self.final_layer_norm(node_features) - - coord_out = self.coord_out(node_features) - - graph_features = scatter(node_features, node2graph, dim=0, reduce="mean") - lattice_out = self.lattice_out(graph_features) - lattice_out = lattice_out.reshape([-1, 3, 3]) - - if self.ip: - lattice_out = paddle.einsum("bij,bjk->bik", lattice_out, lattices) - if self.pred_type: - type_out = self.type_out(node_features) - return lattice_out, coord_out, type_out - - return lattice_out, coord_out diff --git a/ppmat/models/miad/type_diffusion.py b/ppmat/models/miad/type_diffusion.py index 74447cc3..03f8f826 100644 --- a/ppmat/models/miad/type_diffusion.py +++ b/ppmat/models/miad/type_diffusion.py @@ -19,6 +19,7 @@ import paddle import paddle.nn.functional as F +from ppmat.models.miad.diffusion_utils import total_atoms_from_batch from ppmat.models.miad.lattice_diffusion import DDPM from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler @@ -40,8 +41,6 @@ def __init__(self, diffusion_config): "reverse_c0", "reverse_c1", "reverse_std_coef", - "eps_to_x0_c0", - "eps_to_x0_c1", ] for attr in _reshape_attrs: value = getattr(self, attr, None) @@ -55,9 +54,6 @@ def __init__(self, diffusion_config): ).cast("float32") self.from_domain = lambda onehot: onehot.argmax(axis=-1) + 1 - def output_transform(self, x0, batch): - return x0 - def forward_step_sample(self, x0, t, batch): onehot_x0 = self.to_domain(x0) onehot_xt = super().forward_step_sample(onehot_x0, t, batch) @@ -71,8 +67,7 @@ def reverse_step_sample(self, onehot_eps_pred, onehot_xt, t, batch): return onehot_xt_1 def prior_sample(self, batch): - num_atoms = batch["num_atoms"] - total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + total_atoms = total_atoms_from_batch(batch) return paddle.randn([total_atoms, self.num_types], dtype="float32") def loss(self, batch): @@ -84,13 +79,6 @@ def loss(self, batch): ) return l2 - def get_x0_prediction(self, onehot_eps_pred, onehot_xt, t, batch, x0_format="disc"): - onehot_x0_pred = super().get_x0_prediction(onehot_eps_pred, onehot_xt, t, batch) - if x0_format == "onehot": - return onehot_x0_pred - elif x0_format == "disc": - return self.from_domain(onehot_x0_pred) - class D3PM: """Discrete Denoising Diffusion Probabilistic Model for atom types.""" @@ -168,11 +156,6 @@ def _reverse_step_distribution(self, onehot_x0, onehot_xt, t): 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) - xt_1_probs = paddle.where( - paddle.isnan(xt_1_probs), - paddle.full_like(xt_1_probs, 1.0 / self.num_types), - xt_1_probs, - ) if (t.cast("int64") == 0).all(): return self.to_domain(self.from_domain(xt_1_probs.cast("float32"))) xt_1 = ( @@ -184,8 +167,7 @@ def reverse_step_sample(self, onehot_pred, onehot_xt, t, batch): return onehot_xt_1 def prior_sample(self, batch): - num_atoms = batch["num_atoms"] - total_atoms = int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) + total_atoms = total_atoms_from_batch(batch) shape = [total_atoms, self.num_types] xT_probs = paddle.ones(shape, dtype="float32") / self.num_types xT = ( @@ -216,7 +198,3 @@ def loss(self, batch): .sum(axis=-1) ) return self.default_loss_scale * kl_loss - - def get_x0_prediction(self, onehot_pred, onehot_xt, t, batch): - onehot_x0 = self.prediction_to_domain(onehot_pred) - return self.from_domain(onehot_x0) diff --git a/ppmat/schedulers/__init__.py b/ppmat/schedulers/__init__.py index 485c0c84..6da15313 100644 --- a/ppmat/schedulers/__init__.py +++ b/ppmat/schedulers/__init__.py @@ -22,7 +22,6 @@ from ppmat.schedulers.scheduling_lattice_vp import LatticeVPSDEScheduler from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeScheduler from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeSchedulerWrapped -from ppmat.schedulers.miad_schedulers import scheduler as build_miad_scheduler from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal from ppmat.schedulers.scheduling_sde_ve import p_wrapped_normal @@ -31,7 +30,6 @@ ) __all__ = [ "build_scheduler", - "build_miad_scheduler", "DDPMScheduler", "ScoreSdeVeScheduler", "ScoreSdeVeSchedulerWrapped", diff --git a/ppmat/schedulers/miad_schedulers.py b/ppmat/schedulers/miad_schedulers.py index 3f1c5e53..7335257d 100644 --- a/ppmat/schedulers/miad_schedulers.py +++ b/ppmat/schedulers/miad_schedulers.py @@ -26,7 +26,6 @@ import paddle from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal -from ppmat.schedulers.scheduling_sde_ve import p_wrapped_normal # noqa: F401 def scheduler(scheduler_name, num_steps): diff --git a/structure_generation/configs/miad/miad_mp20.yaml b/structure_generation/configs/miad/miad_mp20.yaml index 246c5ec8..e9fca929 100644 --- a/structure_generation/configs/miad/miad_mp20.yaml +++ b/structure_generation/configs/miad/miad_mp20.yaml @@ -86,7 +86,7 @@ Dataset: loader: num_workers: 0 use_shared_memory: False - collate_fn: MiADCollator + collate_fn: DefaultCollator sampler: __class_name__: BatchSampler __init_params__: @@ -108,7 +108,7 @@ Dataset: drop_last: False batch_size: 4 loader: - collate_fn: MiADCollator + collate_fn: DefaultCollator test: dataset: __class_name__: MP20Dataset @@ -124,7 +124,7 @@ Dataset: drop_last: False batch_size: 4 loader: - collate_fn: MiADCollator + collate_fn: DefaultCollator Sample: data: @@ -142,7 +142,7 @@ Sample: drop_last: False batch_size: 4 loader: - collate_fn: MiADCollator + collate_fn: DefaultCollator build_structure_cfg: format: array niggli: False diff --git a/test/test_miad.py b/test/test_miad.py index 2885be8f..454dfebe 100644 --- a/test/test_miad.py +++ b/test/test_miad.py @@ -17,9 +17,10 @@ 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.collate import MiADCollator # noqa: F401 +from ppmat.models.miad.miad import _extract_x0 from ppmat.models.miad.miad import MiAD TINY_MODEL_CFG = { @@ -165,9 +166,10 @@ def test_dataset_sample_fields(self): self.assertIn("num_atoms", sa) def test_collate_fn(self): - collator = MiADCollator() + 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) From 8ec23328bba5cdcea3bb9193d333cdd4f8d113c9 Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Mon, 22 Jun 2026 20:26:30 +0800 Subject: [PATCH 04/10] reuse build_structure --- ppmat/metrics/sun_metric.py | 46 +++++++++++-------------------------- 1 file changed, 13 insertions(+), 33 deletions(-) diff --git a/ppmat/metrics/sun_metric.py b/ppmat/metrics/sun_metric.py index 8e968814..1741db38 100644 --- a/ppmat/metrics/sun_metric.py +++ b/ppmat/metrics/sun_metric.py @@ -38,6 +38,8 @@ from pymatgen.core import Structure from tqdm import tqdm +from ppmat.datasets.build_structure import BuildStructure + logger = logging.getLogger(__name__) @@ -46,35 +48,15 @@ def _composition_hash(structure: Structure) -> str: return str(sorted(list(structure.atomic_numbers))) -def _structures_from_cif_strings(cif_strings: List[str]) -> List[Optional[Structure]]: - """Parse CIF strings into pymatgen Structures.""" - structures = [] - for cif in cif_strings: +def _parse_structures(raw_list, format_str): + """Parse a list of raw data into Structures, returning None on failure.""" + results = [] + for item in raw_list: try: - structures.append(Structure.from_str(cif, fmt="cif")) - except Exception: - structures.append(None) - return structures - - -def _structures_from_dicts(structure_dicts: List[dict]) -> List[Optional[Structure]]: - """Build Structures from dict lists with lattice, atom_types, frac_coords.""" - from pymatgen.core import Lattice - - structures = [] - for d in structure_dicts: - try: - lattice = Lattice(d["lattice"]) - struct = Structure( - lattice, - d["atom_types"], - d["frac_coords"], - coords_are_cartesian=False, - ) - structures.append(struct) + results.append(BuildStructure.build_one(item, format_str, niggli=False, canocial=False)) except Exception: - structures.append(None) - return structures + results.append(None) + return results def compute_uniqueness( @@ -307,9 +289,7 @@ def _load_reference(self): self._reference_structures = [] for cif_str in tqdm(df["cif"], desc="Loading reference structures"): try: - self._reference_structures.append( - Structure.from_str(cif_str, fmt="cif") - ) + 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") @@ -335,16 +315,16 @@ def __call__( # Parse structures if isinstance(generated, pd.DataFrame): if "cif" in generated.columns: - structures = _structures_from_cif_strings(generated["cif"].tolist()) + structures = _parse_structures(generated["cif"].tolist(), "cif_str") elif "structure" in generated.columns: structures = generated["structure"].tolist() else: raise ValueError("DataFrame must have 'cif' or 'structure' column") elif isinstance(generated, list) and len(generated) > 0: if isinstance(generated[0], str): - structures = _structures_from_cif_strings(generated) + structures = _parse_structures(generated, "cif_str") elif isinstance(generated[0], dict): - structures = _structures_from_dicts(generated) + structures = _parse_structures(generated, "array") else: structures = generated else: From 2e914b73fe15e999345700411116a6cef9dcd79b Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Mon, 22 Jun 2026 20:34:54 +0800 Subject: [PATCH 05/10] move miad test file --- test/{ => miad}/test_miad.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename test/{ => miad}/test_miad.py (100%) diff --git a/test/test_miad.py b/test/miad/test_miad.py similarity index 100% rename from test/test_miad.py rename to test/miad/test_miad.py From a4ddf038a45df356f2b88942f2a7355cad69b64a Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Tue, 23 Jun 2026 16:30:17 +0800 Subject: [PATCH 06/10] optimize miad --- ppmat/metrics/__init__.py | 2 +- ppmat/metrics/prerelax_chgnet.py | 83 ---- ppmat/metrics/sun_metric.py | 403 ------------------ ppmat/metrics/sun_metric_utils.py | 294 +++++++++++++ ppmat/models/miad/crystal_diffusion.py | 442 +++++++++++--------- ppmat/models/miad/diffusion_utils.py | 97 ----- ppmat/models/miad/frac_diffusion.py | 154 ------- ppmat/models/miad/lattice_diffusion.py | 212 ---------- ppmat/models/miad/miad.py | 101 +---- ppmat/models/miad/type_diffusion.py | 173 +++----- ppmat/schedulers/__init__.py | 4 - ppmat/schedulers/miad_schedulers.py | 145 ------- structure_generation/configs/miad/README.md | 98 +++-- 13 files changed, 688 insertions(+), 1520 deletions(-) delete mode 100644 ppmat/metrics/prerelax_chgnet.py delete mode 100644 ppmat/metrics/sun_metric.py create mode 100644 ppmat/metrics/sun_metric_utils.py delete mode 100644 ppmat/models/miad/diffusion_utils.py delete mode 100644 ppmat/models/miad/frac_diffusion.py delete mode 100644 ppmat/models/miad/lattice_diffusion.py delete mode 100644 ppmat/schedulers/miad_schedulers.py diff --git a/ppmat/metrics/__init__.py b/ppmat/metrics/__init__.py index 088b1975..0b7d0407 100644 --- a/ppmat/metrics/__init__.py +++ b/ppmat/metrics/__init__.py @@ -18,7 +18,7 @@ from ppmat.metrics.csp_metric import CSPMetric from ppmat.metrics.diffnmr_streaming_adapter import DiffNMRStreamingAdapter -from ppmat.metrics.sun_metric import SUNMetric +from ppmat.metrics.sun_metric_utils import SUNMetric __all__ = [ "build_metric", diff --git a/ppmat/metrics/prerelax_chgnet.py b/ppmat/metrics/prerelax_chgnet.py deleted file mode 100644 index 0e0bcb7a..00000000 --- a/ppmat/metrics/prerelax_chgnet.py +++ /dev/null @@ -1,83 +0,0 @@ -# Copyright (c) 2025 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. - -""" -CHGNet energy prediction utility for S.U.N. Stability metric. -""" - -import logging -from typing import List -from typing import Optional - -import numpy as np -from pymatgen.core import Structure -from tqdm import tqdm - -logger = logging.getLogger(__name__) - -# Lazy-loaded globals to avoid import cost at module level -_chgnet_model = None -_graph_converter = None - - -def _load_chgnet(model_name: str = "chgnet_mptrj"): - """Lazy-load CHGNet model and graph converter.""" - 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", - batch_size: int = 1, -) -> List[Optional[float]]: - """Predict energy per atom for each structure using CHGNet. - - Args: - structures: List of pymatgen Structures (None for invalid). - model_name: CHGNet model name in ppmat registry. - batch_size: Prediction batch size. - - Returns: - List of energy per atom (eV/atom), None for invalid structures. - """ - 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 diff --git a/ppmat/metrics/sun_metric.py b/ppmat/metrics/sun_metric.py deleted file mode 100644 index 1741db38..00000000 --- a/ppmat/metrics/sun_metric.py +++ /dev/null @@ -1,403 +0,0 @@ -# Copyright (c) 2025 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. - -""" -S.U.N. (Stability, Uniqueness, Novelty) metrics for crystal structure generation. - -Implements the evaluation pipeline from MiAD paper for assessing generated -crystal structures across three dimensions: -- Stability: energy above hull below a threshold -- Uniqueness: no duplicates among generated structures -- Novelty: not present in the training set - -Reference: https://arxiv.org/abs/2511.14426 -""" - -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__) - - -def _composition_hash(structure: Structure) -> str: - """Hash a structure by its sorted atomic numbers for quick grouping.""" - return str(sorted(list(structure.atomic_numbers))) - - -def _parse_structures(raw_list, format_str): - """Parse a list of raw data into Structures, returning None on failure.""" - 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 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]: - """Compute uniqueness of generated structures. - - A structure is unique if it does not match any previously seen structure - with the same composition. - - Args: - structures: List of pymatgen Structures (None for invalid). - stol: Site tolerance for StructureMatcher. - angle_tol: Angle tolerance for StructureMatcher. - ltol: Length tolerance for StructureMatcher. - attempt_supercell: Whether to attempt supercell matching. - symmetric: Whether to use symmetric matching. - - Returns: - List of booleans indicating uniqueness. - """ - matcher = StructureMatcher( - stol=stol, angle_tol=angle_tol, ltol=ltol, - attempt_supercell=attempt_supercell, - ) - seen = defaultdict(list) - uniqueness = [] - - for struct in tqdm(structures, desc="Computing uniqueness"): - if struct is None: - uniqueness.append(False) - continue - h = _composition_hash(struct) - if h not in seen: - uniqueness.append(True) - else: - for existing in seen[h]: - if matcher.fit(struct, existing, symmetric=symmetric): - uniqueness.append(False) - break - else: - uniqueness.append(True) - seen[h].append(struct) - - return uniqueness - - -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]: - """Compute novelty of generated structures against a reference set. - - A structure is novel if it does not match any structure in the reference - set (typically the training set). - - Args: - structures: List of pymatgen Structures (None for invalid). - reference_structures: Reference structures (training set). - stol: Site tolerance for StructureMatcher. - angle_tol: Angle tolerance for StructureMatcher. - ltol: Length tolerance for StructureMatcher. - attempt_supercell: Whether to attempt supercell matching. - symmetric: Whether to use symmetric matching. - - Returns: - List of booleans indicating novelty. - """ - matcher = StructureMatcher( - stol=stol, angle_tol=angle_tol, ltol=ltol, - attempt_supercell=attempt_supercell, - ) - - reference_by_comp = defaultdict(list) - for ref in reference_structures: - h = _composition_hash(ref) - reference_by_comp[h].append(ref) - - novelty = [] - for struct in tqdm(structures, desc="Computing novelty"): - if struct is None: - novelty.append(False) - continue - h = _composition_hash(struct) - if h not in reference_by_comp: - novelty.append(True) - else: - for ref in reference_by_comp[h]: - if matcher.fit(struct, ref, symmetric=symmetric): - novelty.append(False) - break - else: - novelty.append(True) - - return novelty - - -def compute_stability( - energy_above_hull: List[Optional[float]], - threshold: float = 0.0, -) -> List[bool]: - """Compute stability from energy above hull values. - - Args: - energy_above_hull: Energy above hull per atom (None for invalid). - threshold: Energy threshold in eV/atom. 0.0 for stable, 0.1 for metastable. - - Returns: - List of booleans indicating stability. - """ - 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]: - """Compute combined S.U.N. metrics. - - S.U.N. = structure is Stable AND Unique AND Novel AND non-trivial - (more than one element type). - - Args: - stability: Stability flags. - uniqueness: Uniqueness flags. - novelty: Novelty flags. - structures: Original structures (for composition check). - min_elements: Minimum number of element types for non-trivial. - - Returns: - Dictionary with metric rates. - """ - n = len(stability) - assert len(uniqueness) == n and len(novelty) == n - - 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: - """S.U.N. metric for evaluating generated crystal structures. - - Computes Stability (energy above hull), Uniqueness (no duplicates), - Novelty (not in training set), and their combination. - - Args: - gt_file_path: Path to ground truth CSV (with 'cif' column) for novelty ref. - stol: Site tolerance for structure matching. - angle_tol: Angle tolerance for structure matching. - ltol: Length tolerance for structure matching. - stability_threshold: Energy above hull threshold in eV/atom. - attempt_supercell: Whether to attempt supercell matching. - """ - - 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, - chgnet_relax: bool = True, - chgnet_relax_steps: int = 200, - ): - 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.chgnet_relax = chgnet_relax - self.chgnet_relax_steps = chgnet_relax_steps - self._reference_structures = None - - def _load_reference(self): - """Load reference structures from ground truth CSV.""" - 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]: - """Run the full S.U.N. evaluation pipeline. - - Args: - generated: Generated structures. Can be: - - List of dicts with 'lattice', 'atom_types', 'frac_coords' - - List of CIF strings - - DataFrame with a 'cif' column - energy_above_hull: Optional pre-computed energy above hull values. - If None, stability is not computed (all marked False). - - Returns: - Dictionary with S.U.N. metrics. - """ - # Parse structures - if isinstance(generated, pd.DataFrame): - if "cif" in generated.columns: - structures = _parse_structures(generated["cif"].tolist(), "cif_str") - elif "structure" in generated.columns: - structures = generated["structure"].tolist() - else: - raise ValueError("DataFrame must have 'cif' or 'structure' column") - elif isinstance(generated, list) and len(generated) > 0: - if isinstance(generated[0], str): - structures = _parse_structures(generated, "cif_str") - elif isinstance(generated[0], dict): - structures = _parse_structures(generated, "array") - else: - structures = generated - else: - raise ValueError("generated must be a list or DataFrame") - - n = len(structures) - logger.info(f"Evaluating {n} generated structures") - - # Uniqueness - uniqueness = compute_uniqueness( - structures, - stol=self.stol, angle_tol=self.angle_tol, ltol=self.ltol, - attempt_supercell=self.attempt_supercell, - ) - - # Novelty - novelty_list = [True] * n # default: all novel - 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, - ) - - # Stability - if energy_above_hull is not None: - stability = compute_stability( - energy_above_hull, threshold=self.stability_threshold - ) - elif self.use_chgnet: - from ppmat.metrics.prerelax_chgnet import predict_energy - - 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 provided and use_chgnet=False, " - "stability metrics will be 0%" - ) - - # Combined S.U.N. - results = compute_sun(stability, uniqueness, novelty_list, structures) - - if self.use_chgnet and energies is not None: - valid_energies = [e for e in energies if e is not None] - if valid_energies: - results["avg_energy_per_atom"] = float(np.mean(valid_energies)) - - # Also compute metastable variant - if energy_above_hull is not None: - metastability = compute_stability( - energy_above_hull, threshold=0.1 - ) - msun_results = compute_sun( - metastability, uniqueness, novelty_list, structures - ) - results["metastable_rate"] = msun_results["stability_rate"] - results["msun_rate"] = msun_results["sun_rate"] - results["msun_count"] = msun_results["sun_count"] - - # Log results - logger.info( - f"S.U.N. Results: " - f"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/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/miad/crystal_diffusion.py b/ppmat/models/miad/crystal_diffusion.py index 4ab55a98..974208b0 100644 --- a/ppmat/models/miad/crystal_diffusion.py +++ b/ppmat/models/miad/crystal_diffusion.py @@ -12,25 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Crystal diffusion controllers for MiAD. -""" - -import os +import math import numpy as np import paddle +from paddle_scatter import scatter_mean from ppmat.models.common.time_embedding import SinusoidalTimeEmbeddings -from ppmat.models.miad.diffusion_utils import TimeDistribution -from ppmat.models.miad.diffusion_utils import mean_interleave -from ppmat.models.miad.frac_diffusion import PFM -from ppmat.models.miad.frac_diffusion import WrappedNormal -from ppmat.models.miad.lattice_diffusion import DDPM -from ppmat.models.miad.lattice_diffusion import FM -from ppmat.models.miad.lattice_diffusion import FM_LenAng from ppmat.models.miad.type_diffusion import D3PM from ppmat.models.miad.type_diffusion import DDPM_onehot +from ppmat.schedulers.scheduling_ddpm import DDPMScheduler +from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeSchedulerWrapped +from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal +from ppmat.utils.crystal import lattice_params_to_matrix_paddle def parse_num_atoms_to_per_crystal(num_atoms_data): @@ -45,51 +39,67 @@ def parse_num_atoms_to_per_crystal(num_atoms_data): return paddle.to_tensor(num_atoms_np.astype("int64")), num_atoms_np - - -def init_diffusion(diffusion_config, logger): - - switch = { - "Default": CrystalGen, - "DiffCSP": CrystalGen, - } - method = diffusion_config.method - if method in switch: - return switch[method](diffusion_config, logger) - raise NotImplementedError(f"Diffusion method '{method}' not implemented") +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: - """Standard crystal generation with single-step reverse sampling.""" + """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.time_distribution = TimeDistribution(self.num_steps, self.cont_time) - time_embed_dim = getattr(self.config, "time_embed_dim", 256) - self.time_embedding = SinusoidalTimeEmbeddings(time_embed_dim) - - # Lattice diffusion - switch_lat = { - "ddpm": DDPM, - "fm": FM, - "fm_lenang": FM_LenAng, - } - self.lat_diffusion = switch_lat[self.config.lat_diffusion.method](self.config) - - # Fractional coordinate diffusion - switch_frac = { - "wrapped_normal": WrappedNormal, - "pfm": PFM, - } - self.frac_diffusion = switch_frac[self.config.frac_diffusion.method]( - self.config + self.eps = 1e-3 + self.time_embedding = SinusoidalTimeEmbeddings( + getattr(self.config, "time_embed_dim", 256) ) - - # Atom type diffusion (only for generation tasks) self.gen_type = "gen" in self.config.task + + # Lattice diffusion + self.lat_method = self.config.lat_diffusion.method + if self.lat_method == "ddpm": + self.lat_scheduler = DDPMScheduler( + num_train_timesteps=self.num_steps, + beta_schedule="squaredcos_cap_v2", + ) + elif 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 + self.frac_method = self.config.frac_diffusion.method + if self.frac_method == "wrapped_normal": + self.frac_scheduler = ScoreSdeVeSchedulerWrapped( + num_train_timesteps=self.num_steps, + sigma_min=0.005, sigma_max=0.5, + sampling_eps=1e-3, + ) + _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, + } + self.step_lr = getattr(self.config.frac_diffusion, "step_lr", None) + if self.step_lr is None: + self.step_lr = _DEFAULT_GAMMA.get(diffusion_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: switch_type = { "ddpm_onehot": DDPM_onehot, @@ -101,191 +111,239 @@ def __init__(self, diffusion_config, logger): 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 - lt = self.lat_diffusion.forward_step_sample(l0, t[0], batch) - ft = self.frac_diffusion.forward_step_sample(f0, t[1], batch) - at = ( - self.type_diffusion.forward_step_sample(a0, t[1], batch) - if self.gen_type - else a0 - ) + 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_diffusion.reverse_step_sample_part_1(f_pred, ft, t[1], 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_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) - ft_1 = self.frac_diffusion.reverse_step_sample_part_2(f_pred, ft_05, t[1], 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_diffusion.reverse_step_sample(l_pred, lt, t[0], batch) - ft_1 = self.frac_diffusion.reverse_step_sample(f_pred, ft, t[1], batch) + 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 + if self.gen_type else at ) return [lt_1, ft_1, at_1] - def _get_batch_info(self, batch): - if "batch" in batch: - return { - "num_atoms": batch["batch"].num_atoms, - "batch_idx": batch["batch"].batch, - "atom_types": batch["batch"].atom_types, - "batch_size": batch["batch_size"], - } - if "batch_idx" in batch: - return { - "num_atoms": batch["num_atoms"], - "batch_idx": batch["batch_idx"], - "atom_types": batch["atom_types"], - "batch_size": batch.get("batch_size", len(batch["num_atoms"])), - } - return self._normalize_batch(batch) - - def _normalize_batch(self, batch): - if "structure_array" in batch: - num_atoms_data = batch["structure_array"].get("num_atoms", None) - else: - num_atoms_data = batch.get("num_atoms", None) - - parsed = parse_num_atoms_to_per_crystal(num_atoms_data) - if parsed is None: - raise ValueError("Cannot determine num_atoms from batch") - 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") - - return { - "num_atoms": num_atoms, - "batch_idx": batch_idx, - "atom_types": atom_types, - "batch_size": batch_size, - } + 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]).reshape([bs, 3]) + ang = 60 + 60 * paddle.rand([4 * 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) + ) + la[:, 3:] = ang[check][: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): - batch_info = self._get_batch_info(batch) - return [ - self.lat_diffusion.prior_sample(batch), - self.frac_diffusion.prior_sample(batch), - ( - self.type_diffusion.prior_sample(batch) - if self.gen_type - else batch_info["atom_types"] - ), - ] + 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): - batch_info = self._get_batch_info(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_info["num_atoms"], batch_info["batch_idx"]) - l_pred = nn_pred[0] - f_pred = nn_pred[1] - a_pred = nn_pred[2] if self.gen_type else None - return [l_pred, f_pred, a_pred] + 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): - modifications = os.environ.get("MODIFICATIONS_FIELD", "") - - batch["t"] = self.time_distribution.sample(batch, 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 - ) - - loss_lat = self.lat_diffusion.loss(batch) - loss_frac = self.frac_diffusion.loss(batch) + batch["prediction"] = self.model_prediction(batch["xt"], batch["t"], model, batch) - if "insert_latloss_" in modifications: - lat_coef = float( - modifications.split("insert_latloss_")[1].split("_coef")[0] - ) - loss_lat = loss_lat * lat_coef - if "insert_fracloss_" in modifications: - frac_coef = float( - modifications.split("insert_fracloss_")[1].split("_coef")[0] - ) - loss_frac = loss_frac * frac_coef + # 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) - loss_lat_val = loss_lat.mean() - loss_frac_val = loss_frac.mean() - batch["loss"] = loss_lat_val + loss_frac_val + batch["loss"] = loss_lat + loss_frac if self.gen_type: - loss_type = self.type_diffusion.loss(batch) - if "insert_typeloss_" in modifications: - type_coef = float( - modifications.split("insert_typeloss_")[1].split("_coef")[0] - ) - loss_type = loss_type * type_coef - loss_type_val = loss_type.mean() - batch["loss"] = batch["loss"] + loss_type_val + loss_type = self.type_diffusion.loss(batch).mean() + batch["loss"] = batch["loss"] + loss_type if self.logger is not None: - batch_info = self._get_batch_info(batch) t_log = paddle.clip(batch["t"][0].clone().cast("int64"), 0, 999) - - self.logger.add( - f"loss:lattice:{mode}", loss_lat_val.item(), stack_after_epoch=True - ) + 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] = loss_lat.clone() - self.logger.add( - f"loss:lattice4time:{mode}", loss_lat_t, stack_after_epoch=True - ) - - self.logger.add( - f"loss:coord:{mode}", loss_frac_val.item(), stack_after_epoch=True - ) + 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( - loss_frac.clone(), batch_info["num_atoms"] - ) - self.logger.add( - f"loss:coord4time:{mode}", loss_frac_t, stack_after_epoch=True + 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_val.item(), - stack_after_epoch=True, - ) - loss_type_t = paddle.zeros([self.num_steps], dtype=loss_type.dtype) - loss_type_t[t_log] = mean_interleave( - loss_type.clone(), batch_info["num_atoms"] - ) - self.logger.add( - f"loss:type4time:{mode}", loss_type_t, stack_after_epoch=True - ) + 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) - reverse_time_iterator = self.time_distribution.reverse_time_iterator( - batch, start_from=self.num_steps - 1 - ) - for t_vector in reverse_time_iterator: - batch["t"] = t_vector - t_value = batch["t"][0][0].item() - progress_printer(t_value) - batch["xt"] = self.reverse_step_sample( - batch["xt"], batch["t"], model, 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 @@ -293,11 +351,7 @@ def sampling_procedure(self, model, batch, progress_printer): def output_transform(self, x0, batch): _out = lambda diff, x: diff.output_transform(x, batch) if hasattr(diff, 'output_transform') else x return [ - _out(self.lat_diffusion, x0[0]), - _out(self.frac_diffusion, x0[1]), - ( - _out(self.type_diffusion, x0[2]) - if self.gen_type - else x0[2] - ), + x0[0], + x0[1], + _out(self.type_diffusion, x0[2]) if self.gen_type else x0[2], ] diff --git a/ppmat/models/miad/diffusion_utils.py b/ppmat/models/miad/diffusion_utils.py deleted file mode 100644 index 3ceb6acc..00000000 --- a/ppmat/models/miad/diffusion_utils.py +++ /dev/null @@ -1,97 +0,0 @@ -# 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. - -""" -Diffusion utilities for MiAD model. -""" - -import os - -import paddle - -from paddle_scatter import scatter_mean - - -class TimeDistribution: - """Time sampling and iteration utilities.""" - - def __init__(self, num_steps, cont_time): - self.num_steps = num_steps - self.cont_time = cont_time - self.eps = 1e-3 - modifications = os.environ.get("MODIFICATIONS_FIELD", "") - if "time_log_normal_0_1_clip" in modifications: - print("MODIF: time_log_normal_0_1_clip", flush=True) - - def _atom_expand(self, batch, t): - if "batch" in batch: - num_atoms = batch["batch"].num_atoms - elif "batch_idx" in batch: - num_atoms = batch["num_atoms"] - else: - raise ValueError("Cannot determine num_atoms from batch") - - t_per_atom = t.repeat_interleave(num_atoms) - return [t, t_per_atom] - - def sample(self, batch, mode=None): - modifications = os.environ.get("MODIFICATIONS_FIELD", "") - if "time_log_normal_0_1_clip" in modifications: - s_eps = 2e-2 - t = paddle.clip( - (1 + 2 * s_eps) - * paddle.nn.functional.sigmoid(paddle.randn(batch["batch_size"])) - - s_eps, - 0, - 1, - ) - else: - t = paddle.rand([batch["batch_size"]]) - - if "max_time_" in modifications: - max_time = float(modifications.split("max_time_")[1].split("+")[0]) - t = t * max_time - - t = self.eps + (self.num_steps - 1 - self.eps) * t - if not self.cont_time: - t = t.round().cast("int64") - - return self._atom_expand(batch, t) - - def reverse_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 self._atom_expand(batch, t) - -def mean_interleave(t, num_repeats): - """Compute per-group mean using paddle_scatter.scatter_mean. - - Args: - t: Tensor of shape [total_atoms] with per-atom values. - num_repeats: Tensor of shape [batch_size] with atom counts per crystal. - - Returns: - Tensor of shape [batch_size] with per-crystal means. - """ - batch_idx = paddle.repeat_interleave( - paddle.arange(num_repeats.shape[0]), num_repeats - ) - return scatter_mean(t, batch_idx, dim=0) - - -def total_atoms_from_batch(batch): - num_atoms = batch["num_atoms"] - return int(num_atoms.sum()) if num_atoms.ndim > 0 else int(num_atoms) diff --git a/ppmat/models/miad/frac_diffusion.py b/ppmat/models/miad/frac_diffusion.py deleted file mode 100644 index 8243028d..00000000 --- a/ppmat/models/miad/frac_diffusion.py +++ /dev/null @@ -1,154 +0,0 @@ -# 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. - -""" -Fractional coordinate diffusion models for MiAD. -""" - -import os - -import paddle - -from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler -from ppmat.models.miad.diffusion_utils import total_atoms_from_batch - - -class WrappedNormal: - """Wrapped Normal diffusion for fractional coordinates.""" - - def __init__(self, diffusion_config): - self.config = diffusion_config.frac_diffusion - self.num_steps = diffusion_config.num_steps - - sigmas_t, sigmas_norm_t, sb, se, d_log_p = get_scheduler( - self.config.scheduler, self.num_steps - ) - - self.sigmas_t = sigmas_t[:, None] - self.sigmas_norm_t = sigmas_norm_t[:, None] - self.sb = sb - self.d_log_p = d_log_p - - _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, - } - self.step_lr = getattr(self.config, "step_lr", None) - if self.step_lr is None: - self.step_lr = _DEFAULT_GAMMA.get(diffusion_config.task, 1e-5) - - self.to_domain = lambda x: x - - def forward_step_sample(self, x0, t, batch): - st = self.sigmas_t[t.cast("int64")] - self.randn_x = paddle.randn(x0.shape) - xt = (x0 + st * self.randn_x) % 1.0 - return xt - - def reverse_step_sample_part_1(self, normed_score_pred, xt, t, batch): - 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 - std_x = paddle.sqrt(2 * step_size) - drift = -step_size * normed_score_pred * paddle.sqrt(snt) - diffusion = std_x * paddle.randn(xt.shape) - xt_05 = xt + drift + diffusion - return xt_05 - - def reverse_step_sample_part_2(self, normed_score_pred, xt_05, t, batch): - 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 - std_x = paddle.sqrt((st_1**2 * (st**2 - st_1**2)) / (st**2)) - drift = -step_size * normed_score_pred * paddle.sqrt(snt) - diffusion = std_x * paddle.randn(xt_05.shape) - xt_1 = (xt_05 + drift + diffusion) % 1.0 - return xt_1 - - def reverse_step_sample(self, normed_score_pred, xt, t, batch): - xt_05 = self.reverse_step_sample_part_1(normed_score_pred, xt, t, batch) - xt_1 = self.reverse_step_sample_part_2(normed_score_pred, xt_05, t, batch) - return xt_1 - - def prior_sample(self, batch): - total_atoms = total_atoms_from_batch(batch) - return paddle.rand([total_atoms, 3], dtype="float32") - - def loss(self, batch): - t_idx = batch["t"][1].cast("int64") - st = self.sigmas_t[t_idx] - snt = self.sigmas_norm_t[t_idx] - normed_score_pred = batch["prediction"][1] - normed_score = self.d_log_p(st * self.randn_x, st) / paddle.sqrt(snt) - l2 = ( - ((normed_score_pred - normed_score) ** 2) - .reshape([normed_score_pred.shape[0], -1]) - .mean(axis=1) - ) - - # mirage infusion code - modifications = os.environ.get("MODIFICATIONS_FIELD", "") - if "miad:add_mirage_atoms_upto" in modifications: - mirage_type = 0 - atom_types = batch["x0"][2] - mask = (atom_types != mirage_type).cast(l2.dtype) - coef = mask.shape[0] / mask.sum() - l2 = l2 * mask * coef - - return l2 - - - -class PFM: - """Periodic Flow Matching for fractional coordinates.""" - - def __init__(self, diffusion_config): - self.config = diffusion_config.frac_diffusion - self.num_steps = diffusion_config.num_steps - self.cont_time = diffusion_config.cont_time - self.final_t = diffusion_config.num_steps - self.step = paddle.to_tensor(1.0 / self.final_t) - self.to_domain = lambda x: x - self.default_loss_scale = 10 - - def forward_step_sample(self, x0, t, batch): - xT_minus_x0 = paddle.rand(x0.shape) - 0.5 - step = (1 + t[:, None]) * self.step - xt = (x0 + step * xT_minus_x0) % 1.0 - self.ut = (-1) * xT_minus_x0 - return xt - - def reverse_step_sample(self, vt, xt, t, batch): - xt_1 = (xt + self.step * vt) % 1.0 - return xt_1 - - def prior_sample(self, batch): - total_atoms = total_atoms_from_batch(batch) - return paddle.rand([total_atoms, 3]) - - def loss(self, batch): - vt = batch["prediction"][1] - l2 = ((vt - self.ut) ** 2).reshape([-1, 3]).mean(axis=1) - return l2 * self.default_loss_scale diff --git a/ppmat/models/miad/lattice_diffusion.py b/ppmat/models/miad/lattice_diffusion.py deleted file mode 100644 index 02e14f8a..00000000 --- a/ppmat/models/miad/lattice_diffusion.py +++ /dev/null @@ -1,212 +0,0 @@ -# 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. - -""" -Lattice diffusion models for MiAD. -""" - -import paddle - -from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler - - -class DDPM: - """DDPM for lattice diffusion. - - Note: This is a lightweight DDPM that pre-computes and stores coefficient - tensors for direct indexing during forward/reverse steps. It differs from - ppmat.schedulers.scheduling_ddpm.DDPMScheduler which is a full-featured - class-based scheduler. The direct-index design is required by MiAD's - CrystalGen which indexes coeff tensors by timestep in tight loops, where - the class-based interface overhead would be prohibitive. - """ - - def __init__(self, diffusion_config, scheduler_cfg=None): - self.config = ( - scheduler_cfg - if scheduler_cfg is not None - else diffusion_config.lat_diffusion - ) - self.num_steps = diffusion_config.num_steps - self.cont_time = diffusion_config.cont_time - - cumprod_alphas_t, ca_t = get_scheduler(self.config.scheduler, self.num_steps) - - cumprod_alphas_t_1 = cumprod_alphas_t[:-1] - cumprod_alphas_t = cumprod_alphas_t[1:] - - if ca_t is not None: - self.f_cumprod_alphas_t = lambda t: ca_t(1 + t).reshape([-1, 1, 1]) - else: - self.f_cumprod_alphas_t = None - - alphas_t = cumprod_alphas_t / cumprod_alphas_t_1 - betas_t = 1 - alphas_t - - # Reshape to (num_steps, 1, 1) for broadcasting with (batch, 3, 3) - self.betas_t = betas_t.reshape([-1, 1, 1]) - self.alphas_t = alphas_t.reshape([-1, 1, 1]) - self.cumprod_alphas_t = cumprod_alphas_t.reshape([-1, 1, 1]) - self.cumprod_alphas_t_1 = cumprod_alphas_t_1.reshape([-1, 1, 1]) - - # Pre-compute reverse sampling coefficients - self.reverse_c0 = 1 / paddle.sqrt(self.alphas_t) - self.reverse_c1 = (1 - self.alphas_t) / paddle.sqrt(1 - self.cumprod_alphas_t) - self.reverse_std_coef = paddle.sqrt( - self.betas_t * (1 - self.cumprod_alphas_t_1) / (1 - self.cumprod_alphas_t) - ) - - self.to_domain = lambda x: x - - def forward_step_sample(self, x0, t, batch): - if self.cont_time and self.f_cumprod_alphas_t is not None: - at = self.f_cumprod_alphas_t(t) - else: - at = self.cumprod_alphas_t[t.cast("int64")] - self.randn_x = paddle.randn(x0.shape) - xt = paddle.sqrt(at) * x0 + paddle.sqrt(1 - at) * self.randn_x - return xt - - def reverse_step_sample(self, eps_pred, xt, t, batch): - t_idx = t.cast("int64") - mu_xt_1 = self.reverse_c0[t_idx] * (xt - self.reverse_c1[t_idx] * eps_pred) - if (t_idx == 0).all(): - return mu_xt_1 - return mu_xt_1 + self.reverse_std_coef[t_idx] * paddle.randn(mu_xt_1.shape) - - def prior_sample(self, batch): - return paddle.randn([batch["batch_size"], 3, 3], dtype="float32") - - def loss(self, batch): - eps_pred = batch["prediction"][0] - l2 = ( - ((eps_pred - self.randn_x) ** 2) - .reshape([eps_pred.shape[0], -1]) - .mean(axis=1) - ) - return l2 - - -class FM: - """Flow Matching for lattice diffusion.""" - - def __init__(self, diffusion_config): - self.config = diffusion_config.lat_diffusion - self.num_steps = diffusion_config.num_steps - self.cont_time = diffusion_config.cont_time - self.step = 1.0 / self.num_steps - self.parameterization = getattr( - diffusion_config.lat_diffusion, "parameterization", "eps" - ) - self.to_domain = lambda x: x - - def forward_step_sample(self, x0, t, batch): - eps = paddle.randn(x0.shape) - step = (1 + t[:, None, None]) * self.step - xt = (1 - step) * x0 + step * eps - if self.parameterization == "eps": - self.ut = (-1) * eps - elif self.parameterization == "v": - self.ut = x0 - eps - return xt - - def reverse_step_sample(self, pred, xt, t, batch): - t_idx = t.cast("int64") - if self.parameterization == "eps": - if t_idx[0] >= 999: - return paddle.randn(xt.shape) - eps_pred = (-1) * pred - step = (1 + t[:, None, None]) * self.step - x0_pred = (xt - step * eps_pred) / (1 - step) - vt = x0_pred - eps_pred - elif self.parameterization == "v": - vt = pred - xt_1 = xt + self.step * vt - return xt_1 - - def prior_sample(self, batch): - return paddle.randn([batch["batch_size"], 3, 3], dtype="float32") - - def loss(self, batch): - vt = batch["prediction"][0] - l2 = ((vt - self.ut) ** 2).reshape([-1, 9]).mean(axis=1) - return l2 - - -class FM_LenAng: - """Flow Matching for lattice in length-angle parameterization.""" - - def __init__(self, diffusion_config): - self.config = diffusion_config.lat_diffusion - self.num_steps = diffusion_config.num_steps - self.cont_time = diffusion_config.cont_time - self.step = 1.0 / self.num_steps - - alpha, theta = 1.3, 0.25 - self.gamma_alpha = alpha - self.gamma_theta = theta - self.angle_difference_bound = 20 - self.default_loss_scale = 0.45 - - self._lenang2lat = None - self.to_domain = lambda x: x - - def _get_converter(self): - if self._lenang2lat is None: - from ppmat.utils.crystal import lattice_params_to_matrix_paddle - - self._lenang2lat = lambda la: lattice_params_to_matrix_paddle( - la[:, :3], la[:, 3:] - ) - return self._lenang2lat - - def forward_step_sample(self, x0, t, batch): - xT = self.prior_sample(batch) - step = (1 + t[:, None, None]) * self.step - xt = (1 - step) * x0 + step * xT - self.ut = x0 - xT - return xt - - def reverse_step_sample(self, vt, xt, t, batch): - xt_1 = xt + self.step * vt - return xt_1 - - def prior_sample(self, batch): - bs = batch["batch_size"] - lenang_xT = paddle.zeros([bs, 6], dtype="float32") - - gamma_dist = paddle.distribution.Gamma( - paddle.to_tensor([self.gamma_alpha], dtype="float32"), - paddle.to_tensor([self.gamma_theta], dtype="float32"), - ) - lenang_xT[:, :3] = 2 + gamma_dist.sample([bs, 3]).reshape([bs, 3]) - - # Sample angles from constrained uniform - ang = 60 + 60 * paddle.rand([4 * 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) - ) - valid = ang[check][:bs] - lenang_xT[:, 3:] = valid - - lenang2lat = self._get_converter() - xT = lenang2lat(lenang_xT) - return xT - - def loss(self, batch): - vt = batch["prediction"][0] - l2 = ((vt - self.ut) ** 2).reshape([-1, 9]).mean(axis=1) - return self.default_loss_scale * l2 diff --git a/ppmat/models/miad/miad.py b/ppmat/models/miad/miad.py index c04781c7..8818c5c3 100644 --- a/ppmat/models/miad/miad.py +++ b/ppmat/models/miad/miad.py @@ -18,10 +18,9 @@ import paddle import paddle.nn as nn -from ppmat.models.miad.crystal_diffusion import init_diffusion +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 import logger from ppmat.utils.crystal import lattices_to_params_shape_numpy @@ -97,101 +96,17 @@ def _gen_edges(num_atoms, frac_coords): self.decoder.gen_edges = _gen_edges if isinstance(diffusion_cfg, dict): diffusion_cfg = _dict_to_sns(diffusion_cfg) - self.diffusion = init_diffusion(diffusion_cfg, logger=None) + self.diffusion = CrystalGen(diffusion_cfg, logger=None) def set_state_dict(self, state_dict, use_structured_name=True): - """ - Load checkpoint with automatic legacy format detection. - - MiAD wraps CSPNet as self.decoder, so all parameter keys in a native - Paddle checkpoint are prefixed with "decoder." - (e.g. decoder.csp_layer_0.edge_mlp.0.weight). - - Legacy PyTorch checkpoints have bare keys (e.g. csp_layer_0.edge_mlp.0.weight) - and Linear weights stored in (out_features, in_features) format. - This method auto-detects and converts them: adds decoder. prefix and transposes - all 2D weight tensors to Paddle's (in_features, out_features) format. - - For legacy checkpoints, ALL 2D Linear weights are unconditionally transposed. - Shape-based detection cannot distinguish square matrices (e.g. 512x512) where - both (out,in) and (in,out) have the same shape. The only safe approach is to - always transpose for legacy-format checkpoints. - """ - model_state = self.state_dict() - model_keys = set(model_state.keys()) - state_keys = set(state_dict.keys()) - - if len(state_keys) == 0: - return super().set_state_dict(state_dict, use_structured_name) - - has_decoder_prefix = any(k.startswith("decoder.") for k in state_keys) - keys_native = state_keys == model_keys or state_keys.issubset(model_keys) - - if keys_native and has_decoder_prefix: - return super().set_state_dict(state_dict, use_structured_name) - - is_legacy = not has_decoder_prefix + 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()} processed = {} - num_transposed = 0 - num_copied = 0 for k, v in state_dict.items(): - new_key = f"decoder.{k}" if is_legacy else k - val_np = v.numpy() if hasattr(v, "numpy") else v - - # Legacy checkpoints always store Linear weights in PyTorch format - # (out_features, in_features). For square matrices (e.g. 512x512), - # shape-based detection cannot work because both orientations have - # identical shapes. Therefore we unconditionally transpose all 2D - # Linear weights when loading from a legacy checkpoint. - if is_legacy and new_key.endswith(".weight") and len(v.shape) == 2: - val_np = val_np.T - num_transposed += 1 - elif is_legacy: - num_copied += 1 - - processed[new_key] = val_np - - model_in_keys = set(model_keys) - set(processed.keys()) - extra_in_ckpt = set(processed.keys()) - set(model_keys) - if model_in_keys: - logger.info( - f"[MiAD] {len(model_in_keys)} model keys missing from checkpoint" - f" (buffers/renamed): {sorted(model_in_keys)[:3]}..." - ) - if extra_in_ckpt: - logger.info( - f"[MiAD] {len(extra_in_ckpt)} checkpoint keys not in model" - f" (skipped): {sorted(extra_in_ckpt)[:3]}..." - ) - if is_legacy: - logger.info( - f"[MiAD] Added decoder. prefix; transposed {num_transposed}" - f" Linear weights, direct copied {num_copied} non-Linear params" - ) - - # Load parameters one by one via set_value to bypass Paddle's - # set_state_dict that may report success without actually loading values. - loaded_missing = [] - loaded_unexpected = [] - for name, param in self.named_parameters(): - if name in processed: - proc_val = processed[name] - if proc_val.shape != tuple(param.shape): - logger.warning( - f"[MiAD] SHAPE MISMATCH: {name}, model={param.shape}," - f" loaded={proc_val.shape} - SKIPPED" - ) - continue - param.set_value(proc_val) - else: - loaded_missing.append(name) - - extra_loaded = set(processed.keys()) - set(model_keys) - if extra_loaded: - loaded_unexpected = list(extra_loaded) - - loaded = (loaded_missing, loaded_unexpected) - return loaded + if k.endswith(".weight") and len(v.shape) == 2: + v = v.T + processed[k] = v + return super().set_state_dict(processed, use_structured_name) def forward(self, batch, **kwargs): mode = "train" if self.training else "val" diff --git a/ppmat/models/miad/type_diffusion.py b/ppmat/models/miad/type_diffusion.py index 03f8f826..372fe388 100644 --- a/ppmat/models/miad/type_diffusion.py +++ b/ppmat/models/miad/type_diffusion.py @@ -12,106 +12,100 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Atom type diffusion models for MiAD. -""" +import math import paddle import paddle.nn.functional as F -from ppmat.models.miad.diffusion_utils import total_atoms_from_batch -from ppmat.models.miad.lattice_diffusion import DDPM -from ppmat.schedulers.miad_schedulers import scheduler as get_scheduler +from ppmat.schedulers.scheduling_ddpm import DDPMScheduler -class DDPM_onehot(DDPM): +class DDPM_onehot: """DDPM for atom types with one-hot encoding.""" def __init__(self, diffusion_config): - super().__init__( - diffusion_config, scheduler_cfg=diffusion_config.type_diffusion - ) - - # Reshape coefficients from (N,1,1) to (N,1) for atom types - _reshape_attrs = [ - "betas_t", - "alphas_t", - "cumprod_alphas_t", - "cumprod_alphas_t_1", - "reverse_c0", - "reverse_c1", - "reverse_std_coef", - ] - for attr in _reshape_attrs: - value = getattr(self, attr, None) - if value is not None and isinstance(value, paddle.Tensor): - if len(value.shape) == 3: - setattr(self, attr, value.reshape([-1, 1])) - + self.num_steps = diffusion_config.num_steps self.num_types = 100 - self.to_domain = lambda types: F.one_hot( - types - 1, num_classes=self.num_types - ).cast("float32") + self.scheduler = DDPMScheduler( + 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) - onehot_xt = super().forward_step_sample(onehot_x0, t, batch) - return onehot_xt + 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, onehot_eps_pred, onehot_xt, t, batch): - onehot_xt_1 = super().reverse_step_sample(onehot_eps_pred, onehot_xt, t, batch) - if (t.cast("int64") == 0).all(): - xt_1 = self.from_domain(onehot_xt_1) - return xt_1 + 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): - total_atoms = total_atoms_from_batch(batch) - return paddle.randn([total_atoms, self.num_types], dtype="float32") + 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] - l2 = ( - ((eps_pred - self.randn_x) ** 2) - .reshape([eps_pred.shape[0], -1]) - .mean(axis=1) - ) - return l2 + return ((eps_pred - self.randn_x) ** 2).reshape([eps_pred.shape[0], -1]).mean(axis=1) class D3PM: """Discrete Denoising Diffusion Probabilistic Model for atom types.""" + @staticmethod + def _uniform_transition_mat(vocab_size, beta_t): + mat = paddle.full((vocab_size, vocab_size), beta_t / float(vocab_size)) + diag_val = 1 - beta_t * (vocab_size - 1) / vocab_size + for i in range(vocab_size): + mat[i, i] = diag_val.cast("float32") + return mat + def __init__(self, diffusion_config): self.config = diffusion_config.type_diffusion self.num_steps = diffusion_config.num_steps self.num_types = 100 - Q_t, cumprod_Q_t = get_scheduler(self.config.scheduler, self.num_steps) + s = 0.008 + discretization = paddle.arange(1, self.num_steps + 1, dtype="float64") + f_t = paddle.cos((discretization / (self.num_steps + 1) + s) / (1 + s) * math.pi / 2) + a_t = f_t / paddle.cos((paddle.to_tensor(0.0, dtype="float64") + s) / (1 + s) * math.pi / 2) + 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(self.num_steps): + Q_t_list.append(self._uniform_transition_mat(self.num_types, betas_t[t_idx])) + Q_t = paddle.stack(Q_t_list, axis=0) + + cumprod_Q_t_list = [Q_t[0]] + for t_idx in range(1, self.num_steps): + 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(Q_t.shape[1], dtype="float32").unsqueeze(0), - Q_t[:-1], - ], - axis=0, + [paddle.eye(Q_t.shape[1], dtype="float32").unsqueeze(0), Q_t[:-1]], axis=0, ) self.Q_t = Q_t.reshape([-1, self.num_types, self.num_types]) self.Q_t_1 = Q_t_1.reshape([-1, self.num_types, self.num_types]) cumprod_Q_t_1 = paddle.concat( - [ - paddle.eye(cumprod_Q_t.shape[1], dtype="float32").unsqueeze(0), - cumprod_Q_t[:-1], - ], - axis=0, + [paddle.eye(cumprod_Q_t.shape[1], dtype="float32").unsqueeze(0), cumprod_Q_t[:-1]], axis=0, ) self.cumprod_Q_t = cumprod_Q_t.reshape([-1, self.num_types, self.num_types]) self.cumprod_Q_t_1 = cumprod_Q_t_1.reshape([-1, self.num_types, self.num_types]) - self.to_domain = lambda types: F.one_hot( - types, num_classes=self.num_types - ).cast("float32") + 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 @@ -121,62 +115,38 @@ def output_transform(self, x0, batch): 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") - ) - onehot_xt = self.to_domain(xt) - return onehot_xt + 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_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 + 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) - result = paddle.where( - paddle.isnan(result), - paddle.full_like(result, 1.0 / self.num_types), - result, - ) - return result + 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") - ) - onehot_xt_1 = self.to_domain(xt_1) - return onehot_xt_1 + 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): - total_atoms = total_atoms_from_batch(batch) - shape = [total_atoms, self.num_types] + 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") - ) - onehot_xT = self.to_domain(xT) - return onehot_xT + 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] @@ -187,14 +157,7 @@ def loss(self, batch): orig_xt_1_probs = self._reverse_step_distribution(onehot_x0, onehot_xt, t) eps = 1e-4 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) + (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/ppmat/schedulers/__init__.py b/ppmat/schedulers/__init__.py index 6da15313..8cd773f1 100644 --- a/ppmat/schedulers/__init__.py +++ b/ppmat/schedulers/__init__.py @@ -22,8 +22,6 @@ from ppmat.schedulers.scheduling_lattice_vp import LatticeVPSDEScheduler from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeScheduler from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeSchedulerWrapped -from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal -from ppmat.schedulers.scheduling_sde_ve import p_wrapped_normal NumAtomsVarianceAdjustedWrappedVESDE = ( scheduling_wrapped_sde_ve.NumAtomsVarianceAdjustedWrappedVESDE @@ -37,8 +35,6 @@ "NumAtomsVarianceAdjustedWrappedVESDE", "D3PMScheduler", "NoiseScheduler", - "p_wrapped_normal", - "d_log_p_wrapped_normal", ] diff --git a/ppmat/schedulers/miad_schedulers.py b/ppmat/schedulers/miad_schedulers.py deleted file mode 100644 index 7335257d..00000000 --- a/ppmat/schedulers/miad_schedulers.py +++ /dev/null @@ -1,145 +0,0 @@ -# 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. - -""" -Pre-computed coefficient schedulers for MiAD-style diffusion models. - -These functions return raw coefficient tensors that diffusion modules index into -during forward/reverse steps. This is a lighter-weight alternative to the class-based -scheduler interface (DDPMScheduler, ScoreSdeVeScheduler, etc.) used by DiffCSP. -""" - -import math - -import numpy as np -import paddle - -from ppmat.schedulers.scheduling_sde_ve import d_log_p_wrapped_normal - - -def scheduler(scheduler_name, num_steps): - """Factory for pre-computed diffusion coefficient tensors. - - Args: - scheduler_name: One of 'diffcsp_cosine', 'cosine', 'default_d3pm', - 'default_wrapped_normal'. - num_steps: Number of diffusion steps. - - Returns: - Tuple of coefficient tensors. The format depends on the scheduler type: - - diffcsp_cosine: (cumprod_alphas_t, None) - - cosine: (cumprod_alphas_t, a_t_fn) - - default_d3pm: (Q_t, cumprod_Q_t) - - default_wrapped_normal: (sigmas, sigmas_norm, sigma_begin, - sigma_end, d_log_p_fn) - """ - if scheduler_name == "diffcsp_cosine": - return _scheduler_diffcsp_cosine(num_steps) - elif scheduler_name == "cosine": - return _scheduler_cosine(num_steps) - elif "default_d3pm" in scheduler_name: - return _scheduler_default_d3pm(num_steps) - elif scheduler_name == "default_wrapped_normal": - return _scheduler_default_wrapped_normal(num_steps) - - raise NotImplementedError(f"Scheduler '{scheduler_name}' not implemented") - - -def _scheduler_diffcsp_cosine(num_steps): - s = 0.008 - discretization = paddle.linspace(0, num_steps, num_steps + 1) - alphas_cumprod = ( - paddle.cos(((discretization / num_steps) + s) / (1 + s) * math.pi * 0.5) ** 2 - ) - alphas_cumprod = alphas_cumprod / alphas_cumprod[0] - betas = 1 - alphas_cumprod[1:] / alphas_cumprod[:-1] - betas = paddle.clip(betas, 0.0001, 0.9999) - alphas = 1.0 - betas - cumprod_alphas_t = paddle.cumprod(alphas, dim=0).cast("float32") - cumprod_alphas_t = paddle.concat([paddle.to_tensor([1.0]), cumprod_alphas_t]) - return cumprod_alphas_t, None - - -def _scheduler_cosine(num_steps): - s = 0.008 - - def f_t(t): - return paddle.cos((t / (num_steps + 1) + s) / (1 + s) * math.pi / 2) ** 2 - - def a_t(t): - return f_t(t) / f_t(paddle.to_tensor(0.0, dtype="float64")) - - discretization = paddle.linspace(0, num_steps, num_steps + 1) - cumprod_alphas_t = a_t(discretization).cast("float32") - return cumprod_alphas_t, a_t - - -def _scheduler_default_d3pm(num_steps, vocab_size=100): - def get_uniform_transition_mat(vocab_size, beta_t): - mat = paddle.full((vocab_size, vocab_size), beta_t / float(vocab_size)) - diag_val = 1 - beta_t * (vocab_size - 1) / vocab_size - for i in range(vocab_size): - mat[i, i] = diag_val.cast("float32") - return mat - - s = 0.008 - - def f_t(t): - return paddle.cos((t / (num_steps + 1) + s) / (1 + s) * math.pi / 2) - - def a_t(t): - return f_t(t) / f_t(paddle.to_tensor(0.0, dtype="float64")) - - discretization = paddle.arange(1, num_steps + 1, dtype="float64") - cumprod_alphas_t = a_t(discretization) - 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_steps): - Q_t_list.append(get_uniform_transition_mat(vocab_size, betas_t[t_idx])) - Q_t = paddle.stack(Q_t_list, axis=0) - - cumprod_Q_t_list = [Q_t[0]] - for t_idx in range(1, num_steps): - 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) - - return Q_t, cumprod_Q_t - - -def _scheduler_default_wrapped_normal(num_steps): - """Wrapped normal scheduler for fractional coordinate diffusion. - - Uses sn=100000 for sigma_norm calculation (differs from scheduling_sde_ve.py - which uses sn=10000). This larger sample size was used in the original MiAD - implementation for more precise expected score norm estimation. - """ - sigma_begin, sigma_end = 0.005, 0.5 - sigmas = paddle.to_tensor( - np.exp(np.linspace(np.log(sigma_begin), np.log(sigma_end), num_steps)), - dtype="float32", - ) - - def sigma_norm(sigma, T=1.0, sn=100000): - sigmas_expanded = paddle.tile(sigma[None, :], [sn, 1]) - x_sample = sigma * paddle.randn(sigmas_expanded.shape) - x_sample = x_sample % T - normal_ = d_log_p_wrapped_normal(x_sample, sigmas_expanded, T=T) - return (normal_**2).mean(axis=0) - - sigmas_norm_ = sigma_norm(sigmas) - return sigmas, sigmas_norm_, sigma_begin, sigma_end, d_log_p_wrapped_normal diff --git a/structure_generation/configs/miad/README.md b/structure_generation/configs/miad/README.md index 1bd0938e..359b0db0 100644 --- a/structure_generation/configs/miad/README.md +++ b/structure_generation/configs/miad/README.md @@ -1,45 +1,70 @@ # MiAD -MiAD (Mirage Atom Diffusion) is a diffusion-based framework for de novo crystal generation. It introduces the concept of 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. +[Mirage Atom Diffusion for De Novo Crystal Generation](https://arxiv.org/abs/2511.14426) -## How It Works +## Abstract -1. Mirage Atoms: The model uses sentinel nodes called "mirage atoms" (assigned atom_type = 0) that can be transformed into real chemical elements or discarded as empty space during the diffusion trajectory . +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. -2. Training Phase: Real crystals are padded with mirage atoms up to a maximum limit (e.g., 25 atoms), where mirage nodes have random fractional coordinates and type 0 . +![MiAD Overview](https://github.com/andrey-okhotin/miad/blob/main/pictures_from_paper/miad_method_scheme.png) -3. Generation Phase: The model initializes all crystals with the maximum atom count and predicts which nodes should materialize into real atoms versus remaining as mirage atoms . +--- -4. Post-Processing: Remaining mirage atoms are stripped before exporting to final CIF files . +## Model Description -## Architecture +### Overview -MiAD uses DiffCSP as its backbone architecture, with the CrystalGen orchestrator handling the diffusion process for lattice parameters, fractional coordinates, and atom types. +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}$ -## Dataset +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) +$$ -[Original dataset address](https://drive.google.com/file/d/1BLI3VtvzfIIXlH6UHQ4o-gQaCIOZ1UR7/view?usp=sharing) +Also supports Periodic Flow Matching. -[Mirror AIStudio address](https://aistudio.baidu.com/modelsdetail/48578/intro) +#### 3) Atom-type diffusion +Atom types use D3PM (uniform transition + cosine schedule) or DDPM with one-hot encoding. -Extract the data to `./data/mp_20/` so that the CSV files are at `./data/mp_20/train.csv`, `./data/mp_20/val.csv`, and `./data/mp_20/test.csv`. +#### 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. -## Environment Dependencies +--- -- Python >= 3.10 -- PaddlePaddle >= 3.3.0 (official release) -- paddle_scatter >= 2.1.2 -- numpy, pymatgen, omegaconf +## Dataset -## Checkpoints +| Source | Link | +|-----------------------|------| +| Original Google Drive | [Google Drive](https://drive.google.com/file/d/1BLI3VtvzfIIXlH6UHQ4o-gQaCIOZ1UR7/view?usp=sharing) | +| Mirror AiStudio | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) | -[Original checkpoints address](https://drive.google.com/file/d/1KyD6KzvjYFPfU8lutFyO_0b8EbeHSGqf/view?usp=sharing) +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`. -[Mirror AIStudio address](https://aistudio.baidu.com/modelsdetail/48578/intro) +--- -[Paddle checkpoints address](https://aistudio.baidu.com/modelsdetail/48638/intro) +## Results +| Metric | CHGNet | eq-V2 | +|--------|--------|-------| +| S.U.N. | 12.21% | 5.64% | +--- ## Commands @@ -59,19 +84,34 @@ python -m paddle.distributed.launch --gpus="0,1,2,3" structure_generation/train. 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' ``` -### Sampling (Structure Generation) +### Sampling ```bash -# Option 1: Use pre-trained model (auto-download) +# Mode 1: pre-trained model (auto-download) 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 -# Option 2: Custom checkpoint +# Mode 2: custom checkpoint 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' ``` -### Evaluation (Compute Metrics) +### Checkpoints -```bash -# Evaluate generated structures against ground truth using CSPMetric -python structure_generation/sample.py --model_name='miad_mp20' --weights_name='latest.pdparams' --save_path='result_eval/' --mode='compute_metric' -``` \ No newline at end of file +| Source | Link | +|--------------------------------|------| +| Original (PyTorch format) | [Google Drive](https://drive.google.com/file/d/1KyD6KzvjYFPfU8lutFyO_0b8EbeHSGqf/view?usp=sharing) | +| Mirror (PyTorch format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) | +| PaddleMaterials(Paddle format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48638/intro) | + +--- + +## 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} +} +``` From 4118fa9c7713a5b338e00f6ace0ddf4a45418966 Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Sat, 27 Jun 2026 00:55:26 +0800 Subject: [PATCH 07/10] optimize miad --- ppmat/models/diffcsp/diffcsp.py | 18 ++++--- ppmat/models/miad/crystal_diffusion.py | 40 +++++++------- ppmat/models/miad/miad.py | 72 +++++++++++++++----------- ppmat/models/miad/type_diffusion.py | 4 +- 4 files changed, 77 insertions(+), 57 deletions(-) diff --git a/ppmat/models/diffcsp/diffcsp.py b/ppmat/models/diffcsp/diffcsp.py index 05721eb1..8d565ba7 100644 --- a/ppmat/models/diffcsp/diffcsp.py +++ b/ppmat/models/diffcsp/diffcsp.py @@ -66,6 +66,7 @@ def __init__( dis_emb=None, ln=False, ip=True, + use_prop_mlp=True, ): super(CSPLayer, self).__init__() self.dis_dim = 3 @@ -88,12 +89,13 @@ def __init__( act_fn, ) - self.prop_mlp = paddle.nn.Sequential( - paddle.nn.Linear(in_features=prop_dim, out_features=hidden_dim), - act_fn, - paddle.nn.Linear(in_features=hidden_dim, out_features=hidden_dim), - act_fn, - ) + if use_prop_mlp: + self.prop_mlp = paddle.nn.Sequential( + paddle.nn.Linear(in_features=prop_dim, out_features=hidden_dim), + act_fn, + paddle.nn.Linear(in_features=hidden_dim, out_features=hidden_dim), + act_fn, + ) self.ln = ln if self.ln: @@ -149,7 +151,7 @@ def forward( property_emb=None, property_mask=None, ): - if property_emb is not None: + if property_emb is not None and hasattr(self, "prop_mlp"): property_features = self.prop_mlp(property_emb) if property_mask is not None: property_features = property_features * property_mask @@ -213,6 +215,7 @@ def __init__( prop_dim: int = 512, pred_scalar: bool = False, num_classes: int = 100, + use_prop_mlp: bool = True, ): super(CSPNet, self).__init__() self.ip = ip @@ -247,6 +250,7 @@ def __init__( dis_emb=self.dis_emb, ln=ln, ip=ip, + use_prop_mlp=use_prop_mlp, ), ) self.num_layers = num_layers diff --git a/ppmat/models/miad/crystal_diffusion.py b/ppmat/models/miad/crystal_diffusion.py index 974208b0..e295d4cd 100644 --- a/ppmat/models/miad/crystal_diffusion.py +++ b/ppmat/models/miad/crystal_diffusion.py @@ -12,20 +12,26 @@ # See the License for the specific language governing permissions and # limitations under the License. -import math - 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 DDPM_onehot +from ppmat.models.miad.type_diffusion import DDPMOnehot from ppmat.schedulers.scheduling_ddpm import DDPMScheduler from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeSchedulerWrapped 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: @@ -85,13 +91,6 @@ def __init__(self, diffusion_config, logger): sigma_min=0.005, sigma_max=0.5, sampling_eps=1e-3, ) - _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, - } self.step_lr = getattr(self.config.frac_diffusion, "step_lr", None) if self.step_lr is None: self.step_lr = _DEFAULT_GAMMA.get(diffusion_config.task, 1e-5) @@ -102,7 +101,7 @@ def __init__(self, diffusion_config, logger): # Type diffusion if self.gen_type: switch_type = { - "ddpm_onehot": DDPM_onehot, + "ddpm_onehot": DDPMOnehot, "d3pm": D3PM, } self.type_diffusion = switch_type[self.config.type_diffusion.method]( @@ -253,14 +252,17 @@ def _prior_lat(self, batch): paddle.to_tensor([self.gamma_alpha], dtype="float32"), paddle.to_tensor([self.gamma_theta], dtype="float32"), ) - la[:, :3] = 2 + gamma.sample([bs, 3]).reshape([bs, 3]) - ang = 60 + 60 * paddle.rand([4 * 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) - ) - la[:, 3:] = ang[check][:bs] + 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}") diff --git a/ppmat/models/miad/miad.py b/ppmat/models/miad/miad.py index 8818c5c3..f7191779 100644 --- a/ppmat/models/miad/miad.py +++ b/ppmat/models/miad/miad.py @@ -24,6 +24,16 @@ from ppmat.utils.crystal import lattices_to_params_shape_numpy +class MiADCSPNet(CSPNet): + """CSPNet with block-diagonal edge generation for GPU compatibility.""" + + 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]]) % 1.0 + + def _to_numpy(x): if hasattr(x, "numpy"): return x.numpy() @@ -81,19 +91,7 @@ def __init__(self, model_cfg=None, diffusion_cfg=None, **kwargs): "prop_dim", "pred_scalar", "num_classes", } cspnet_kwargs = {k: v for k, v in model_cfg.items() if k in _cspnet_keys} - self.decoder = CSPNet(**cspnet_kwargs) - # Remove unused prop_mlp (parent creates it; checkpoint lacks these weights) - for i in range(self.decoder.num_layers): - layer = getattr(self.decoder, f"csp_layer_{i}", None) - if layer and hasattr(layer, "prop_mlp"): - del layer.prop_mlp - # Replace gen_edges with block_diag (parent meshgrid causes GPU crash on certain batch sizes) - def _gen_edges(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]]) - self.decoder.gen_edges = _gen_edges + self.decoder = MiADCSPNet(use_prop_mlp=False, **cspnet_kwargs) if isinstance(diffusion_cfg, dict): diffusion_cfg = _dict_to_sns(diffusion_cfg) self.diffusion = CrystalGen(diffusion_cfg, logger=None) @@ -101,12 +99,26 @@ def _gen_edges(num_atoms, frac_coords): 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()} - processed = {} + expected_shapes = { + k: v.shape for k, v in self.state_dict().items() + } for k, v in state_dict.items(): - if k.endswith(".weight") and len(v.shape) == 2: - v = v.T - processed[k] = v - return super().set_state_dict(processed, use_structured_name) + expected = expected_shapes.get(k) + if expected is not None and v.shape != expected and v.T.shape == expected: + state_dict[k] = v.T + 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 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" @@ -142,6 +154,7 @@ def sample(self, batch_data, num_inference_steps=None): 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, @@ -157,14 +170,15 @@ def sample(self, batch_data, num_inference_steps=None): def progress_printer(t): return None - batch = self.diffusion.sampling_procedure( - model=self.decoder, - batch=batch_data, - progress_printer=progress_printer, - ) - - if original_steps is not None: - self.diffusion.num_steps = original_steps + 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"] @@ -190,9 +204,9 @@ def progress_printer(t): lat_i = lattices[i] fc_i = frac_coords[start_idx : start_idx + n] at_i = atom_types[start_idx : start_idx + n] - lat_np = lat_i.numpy() if hasattr(lat_i, "numpy") else np.asarray(lat_i) - fc_np = fc_i.numpy() if hasattr(fc_i, "numpy") else np.asarray(fc_i) - at_np = at_i.numpy() if hasattr(at_i, "numpy") else np.asarray(at_i) + 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. diff --git a/ppmat/models/miad/type_diffusion.py b/ppmat/models/miad/type_diffusion.py index 372fe388..bb9d8efb 100644 --- a/ppmat/models/miad/type_diffusion.py +++ b/ppmat/models/miad/type_diffusion.py @@ -20,7 +20,7 @@ from ppmat.schedulers.scheduling_ddpm import DDPMScheduler -class DDPM_onehot: +class DDPMOnehot: """DDPM for atom types with one-hot encoding.""" def __init__(self, diffusion_config): @@ -155,7 +155,7 @@ def loss(self, batch): 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-4 + 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) From 8b4136d4aa707c708324579e47817f11729c14ca Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Sat, 27 Jun 2026 14:40:44 +0800 Subject: [PATCH 08/10] optimize miad --- ppmat/models/diffcsp/diffcsp.py | 18 +++++++----------- ppmat/models/miad/miad.py | 17 +++++++---------- 2 files changed, 14 insertions(+), 21 deletions(-) diff --git a/ppmat/models/diffcsp/diffcsp.py b/ppmat/models/diffcsp/diffcsp.py index 8d565ba7..05721eb1 100644 --- a/ppmat/models/diffcsp/diffcsp.py +++ b/ppmat/models/diffcsp/diffcsp.py @@ -66,7 +66,6 @@ def __init__( dis_emb=None, ln=False, ip=True, - use_prop_mlp=True, ): super(CSPLayer, self).__init__() self.dis_dim = 3 @@ -89,13 +88,12 @@ def __init__( act_fn, ) - if use_prop_mlp: - self.prop_mlp = paddle.nn.Sequential( - paddle.nn.Linear(in_features=prop_dim, out_features=hidden_dim), - act_fn, - paddle.nn.Linear(in_features=hidden_dim, out_features=hidden_dim), - act_fn, - ) + self.prop_mlp = paddle.nn.Sequential( + paddle.nn.Linear(in_features=prop_dim, out_features=hidden_dim), + act_fn, + paddle.nn.Linear(in_features=hidden_dim, out_features=hidden_dim), + act_fn, + ) self.ln = ln if self.ln: @@ -151,7 +149,7 @@ def forward( property_emb=None, property_mask=None, ): - if property_emb is not None and hasattr(self, "prop_mlp"): + if property_emb is not None: property_features = self.prop_mlp(property_emb) if property_mask is not None: property_features = property_features * property_mask @@ -215,7 +213,6 @@ def __init__( prop_dim: int = 512, pred_scalar: bool = False, num_classes: int = 100, - use_prop_mlp: bool = True, ): super(CSPNet, self).__init__() self.ip = ip @@ -250,7 +247,6 @@ def __init__( dis_emb=self.dis_emb, ln=ln, ip=ip, - use_prop_mlp=use_prop_mlp, ), ) self.num_layers = num_layers diff --git a/ppmat/models/miad/miad.py b/ppmat/models/miad/miad.py index f7191779..d0f22c2c 100644 --- a/ppmat/models/miad/miad.py +++ b/ppmat/models/miad/miad.py @@ -25,13 +25,15 @@ class MiADCSPNet(CSPNet): - """CSPNet with block-diagonal edge generation for GPU compatibility.""" + """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]]) % 1.0 + return fc_edges, frac_coords[fc_edges[1]] - frac_coords[fc_edges[0]] def _to_numpy(x): @@ -91,7 +93,7 @@ def __init__(self, model_cfg=None, diffusion_cfg=None, **kwargs): "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(use_prop_mlp=False, **cspnet_kwargs) + self.decoder = MiADCSPNet(**cspnet_kwargs) if isinstance(diffusion_cfg, dict): diffusion_cfg = _dict_to_sns(diffusion_cfg) self.diffusion = CrystalGen(diffusion_cfg, logger=None) @@ -99,13 +101,6 @@ def __init__(self, model_cfg=None, diffusion_cfg=None, **kwargs): 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()} - expected_shapes = { - k: v.shape for k, v in self.state_dict().items() - } - for k, v in state_dict.items(): - expected = expected_shapes.get(k) - if expected is not None and v.shape != expected and v.T.shape == expected: - state_dict[k] = v.T param_state = {} for name, param in self.named_parameters(): if name in state_dict: @@ -114,6 +109,8 @@ def set_state_dict(self, state_dict, use_structured_name=True): 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: From 38036f1f6041023d9b8bcf10c35859193707c0af Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Sat, 27 Jun 2026 14:43:07 +0800 Subject: [PATCH 09/10] optimize miad --- README.md | 62 +++++++++++++++++++++++++++---------------------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index 3f2491bf..dc0adba4 100755 --- a/README.md +++ b/README.md @@ -12,29 +12,29 @@

-### Core Capabilities +### 🧩 Core Capabilities | Task | Description | Typical Applications | |------|-------------|---------------------| -| **Property Prediction (PP)** | Predict material properties from structure | Formation energy, band gap, elastic moduli | -| **Structure Generation (SG)** | Generate novel crystal structures | High-throughput screening, inverse design | -| **Interatomic Potential (IP)** | Replace DFT with ML potentials | Molecular dynamics, large-scale simulations | -| **Electronic Structure (ES)** | Predict electronic properties | Band structure, density of states | +| **Property Prediction (PP)** | Predict material properties from structure | Forward design or predict formation energy, band gap, elastic moduli etc. | +| **Structure Generation (SG)** | Generate novel crystal structures | Inverse design or structure generation | +| **Machine Learning Interatomic Potential (MLIP)** | Surrogate Model for DFT as ML potentials | Molecular dynamics simulations | +| **Electronic Structure (ES)** | Surrogate Model for DFT to predict physical field | Predict electronic density | | **Spectrum Elucidation (SE)** | Reconstruct structures from spectra | NMR structure elucidation | | **Spectrum Enhancement (SPEN)** | Enhance microscopy and spectrum signals | STEM image enhancement, denoising | -### Supported Materials +### 🧱 Supported Materials - **Inorganic Crystals** - Well-supported with multiple datasets and pretrained models - **Organic Molecules** - Support for small molecule datasets and property prediction - *Polymers, catalysts, and amorphous materials are under development* -### Why PaddleMaterials? +### ✨ Why PaddleMaterials? -- ✅ **Rich Pretrained Models** - 50+ pretrained models ready for inference -- ✅ **Multi-Task Integration** - Unified framework across PP, SG, MLIP, MLES, SE, SPEN -- ✅ **Domestic Hardware Support** - Full support for MetaX GPUs and NVIDIA GPUs -- ✅ **Production-Ready** - Distributed training, mixed precision, checkpoint recovery +- ✅ **Rich Pretrained Models & AI-ready Datasets** - 50+ pretrained models ready for inference and Multiple curated datasets for training +- ✅ **Multi-Task Integration** - Unified framework across tasks of PP, SG, MLIP, ES, SE, SPEN etc. +- ✅ **Multi-Hardware Support** - Full support for NVIDIA GPUs and MetaX GPUs and Intel CPUs +- ✅ **Production-Ready** - Easy to use with standandlize design & distributed training, mixed precision, checkpoint recovery ### 📑 Support Tasks @@ -42,27 +42,27 @@ |------|-------------|------| | **Property Prediction (PP)** | Predict formation energy, band gap, elastic properties | [README](property_prediction/README.md) | | **Structure Generation (SG)** | Generate new crystal structures with diffusion models | [README](structure_generation/README.md) | -| **Interatomic Potential (IP)** | DFT-accurate potentials for molecular dynamics | [README](interatomic_potentials/README.md) | +| **Machine Learning Interatomic Potential (MLIP)** | DFT-accurate potentials for molecular dynamics | [README](interatomic_potentials/README.md) | | **Electronic Structure (ES)** | Predict electronic structure properties | [README](electronic_structure/README.md) | | **Spectrum Elucidation (SE)** | Reconstruct molecular structures from NMR spectra | [README](spectrum_elucidation/README.md) | -| **Spectrum Enhancement (SPEN)** | Enhance microscopy and spectrum signals | [README](spectrum_enhancement/README.md) | +| **Spectrum Enhancement (SPEN)** | Enhance microscopy and spectral signals | [README](spectrum_enhancement/README.md) | -### 🎯 Available Pretrained Models +### 🤖 Available Pretrained Models | Task | Models | Dataset | |------|--------|---------| -| **Property Prediction** | MEGNet, iComformer, DimeNet++ | MP2018, MP2024, JARVIS | -| **Structure Generation** | MatterGen, DiffCSP | MP20, ALEX | -| **Interatomic Potentials** | CHGNet, MatterSim | MPTRJ | -| **Electronic Structure** | InfGCN | Custom datasets | -| **Spectrum Elucidation** | DiffNMR | MSD_NMR | -| **Spectrum Enhancement** | SFIN | HAADF/BF STEM image datasets | +| **Property Prediction** | MEGNet, iComformer, DimeNet++ | MP2018, MP2024, JARVIS | +| **Structure Generation** | MatterGen, DiffCSP | MP20, ALEX | +| **Machine Learning Interatomic Potential** | CHGNet, MatterSim | MPTRJ | +| **Electronic Structure** | InfGCN | QM9_ES, MP_ES, OMol25_MC_ES | +| **Spectrum Elucidation** | DiffNMR | MSD_NMR | +| **Spectrum Enhancement** | SFIN | SFIN-HAADF/BF | Full model list: See [MODEL_REGISTRY](ppmat/models/__init__.py#L75) --- -## ⚡ Get Started +## 🚀 Get Started ### 🔧 Installation @@ -70,7 +70,7 @@ Please refer to the installation [document](Install.md) for your hardware enviro --- -### Mimi Inference +### ⚡ Easy Inference #### Property Prediction @@ -141,13 +141,15 @@ python spectrum_enhancement/predict.py --- -### Train Your Own Model +### 🏋️ Start Training For training and fine-tuning, refer to the [documentation](get_started.md). --- -## 👩‍👩‍👧‍👦 Contributors & Cooperation & Community +## 🤝 Contributors & Cooperation & Community + +[![Star History Chart](https://api.star-history.com/svg?repos=PaddlePaddle/PaddleMaterials&type=date&legend=top-left)](https://www.star-history.com/#PaddlePaddle/PaddleMaterilas&type=date&legend=top-left) Thanks to all contributors who have helped build PaddleMaterials！ @@ -156,19 +158,17 @@ Thanks to all contributors who have helped build PaddleMaterials！ Thanks for the following organiziton for cooprative support!

- - - + + +

-[![Star History Chart](https://api.star-history.com/svg?repos=PaddlePaddle/PaddleMaterials&type=date&legend=top-left)](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: From 62aa181c469ab16952b6f1993c04627719a11750 Mon Sep 17 00:00:00 2001 From: caoyuanye Date: Tue, 30 Jun 2026 16:02:35 +0800 Subject: [PATCH 10/10] optimize pr review --- ppmat/models/miad/crystal_diffusion.py | 65 ++++++----- ppmat/models/miad/miad.py | 11 -- ppmat/models/miad/type_diffusion.py | 107 ++++++++++++------ structure_generation/configs/miad/README.md | 66 ++++++----- .../configs/miad/miad_mp20.yaml | 15 ++- test/miad/__init__.py | 13 +++ test/miad/test_miad.py | 26 ++++- 7 files changed, 198 insertions(+), 105 deletions(-) create mode 100644 test/miad/__init__.py diff --git a/ppmat/models/miad/crystal_diffusion.py b/ppmat/models/miad/crystal_diffusion.py index e295d4cd..319628d0 100644 --- a/ppmat/models/miad/crystal_diffusion.py +++ b/ppmat/models/miad/crystal_diffusion.py @@ -19,8 +19,7 @@ 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.scheduling_ddpm import DDPMScheduler -from ppmat.schedulers.scheduling_sde_ve import ScoreSdeVeSchedulerWrapped +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 @@ -63,50 +62,56 @@ class CrystalGen: 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.cont_time = self.config["cont_time"] + self.num_steps = self.config["num_steps"] self.eps = 1e-3 self.time_embedding = SinusoidalTimeEmbeddings( - getattr(self.config, "time_embed_dim", 256) + self.config.get("time_embed_dim", 256) ) - self.gen_type = "gen" in self.config.task + self.gen_type = "gen" in self.config["task"] # Lattice diffusion - self.lat_method = self.config.lat_diffusion.method - if self.lat_method == "ddpm": - self.lat_scheduler = DDPMScheduler( - num_train_timesteps=self.num_steps, - beta_schedule="squaredcos_cap_v2", - ) - elif self.lat_method == "fm_lenang": + 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 - self.frac_method = self.config.frac_diffusion.method - if self.frac_method == "wrapped_normal": - self.frac_scheduler = ScoreSdeVeSchedulerWrapped( - num_train_timesteps=self.num_steps, - sigma_min=0.005, sigma_max=0.5, - sampling_eps=1e-3, - ) - self.step_lr = getattr(self.config.frac_diffusion, "step_lr", None) - if self.step_lr is None: - self.step_lr = _DEFAULT_GAMMA.get(diffusion_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 + 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[self.config.type_diffusion.method]( - self.config - ) + self.type_diffusion = switch_type[type_cfg["method"]](self.config) else: self.type_diffusion = None @@ -179,7 +184,7 @@ def forward_step_sample(self, x0, t, batch): def reverse_step_sample(self, xt, t, model, batch): lt, ft, at = xt - if self.config.method == "DiffCSP": + 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] diff --git a/ppmat/models/miad/miad.py b/ppmat/models/miad/miad.py index d0f22c2c..b7d620d4 100644 --- a/ppmat/models/miad/miad.py +++ b/ppmat/models/miad/miad.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from types import SimpleNamespace - import numpy as np import paddle import paddle.nn as nn @@ -69,13 +67,6 @@ def _extract_x0(batch): return batch -def _dict_to_sns(d): - """Recursively convert dict to SimpleNamespace for attribute access.""" - if not isinstance(d, dict): - return d - return SimpleNamespace(**{k: _dict_to_sns(v) for k, v in d.items()}) - - class MiAD(nn.Layer): """Mirage Atom Diffusion model.""" @@ -94,8 +85,6 @@ def __init__(self, model_cfg=None, diffusion_cfg=None, **kwargs): } cspnet_kwargs = {k: v for k, v in model_cfg.items() if k in _cspnet_keys} self.decoder = MiADCSPNet(**cspnet_kwargs) - if isinstance(diffusion_cfg, dict): - diffusion_cfg = _dict_to_sns(diffusion_cfg) self.diffusion = CrystalGen(diffusion_cfg, logger=None) def set_state_dict(self, state_dict, use_structured_name=True): diff --git a/ppmat/models/miad/type_diffusion.py b/ppmat/models/miad/type_diffusion.py index bb9d8efb..37b83f3e 100644 --- a/ppmat/models/miad/type_diffusion.py +++ b/ppmat/models/miad/type_diffusion.py @@ -17,18 +17,23 @@ import paddle import paddle.nn.functional as F -from ppmat.schedulers.scheduling_ddpm import DDPMScheduler +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_steps = diffusion_config["num_steps"] self.num_types = 100 - self.scheduler = DDPMScheduler( - num_train_timesteps=self.num_steps, - beta_schedule="squaredcos_cap_v2", + 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 @@ -58,26 +63,32 @@ def loss(self, batch): return ((eps_pred - self.randn_x) ** 2).reshape([eps_pred.shape[0], -1]).mean(axis=1) -class D3PM: - """Discrete Denoising Diffusion Probabilistic Model for atom types.""" +class D3PMUniformScheduler: + """D3PM scheduler with uniform transition matrices (model-internal). - @staticmethod - def _uniform_transition_mat(vocab_size, beta_t): - mat = paddle.full((vocab_size, vocab_size), beta_t / float(vocab_size)) - diag_val = 1 - beta_t * (vocab_size - 1) / vocab_size - for i in range(vocab_size): - mat[i, i] = diag_val.cast("float32") - return mat + 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, diffusion_config): - self.config = diffusion_config.type_diffusion - self.num_steps = diffusion_config.num_steps - self.num_types = 100 + def __init__( + self, + num_train_timesteps: int = 1000, + num_types: int = 100, + s: float = 0.008, + ): + self.num_types = num_types - s = 0.008 - discretization = paddle.arange(1, self.num_steps + 1, dtype="float64") - f_t = paddle.cos((discretization / (self.num_steps + 1) + s) / (1 + s) * math.pi / 2) - a_t = f_t / paddle.cos((paddle.to_tensor(0.0, dtype="float64") + s) / (1 + s) * math.pi / 2) + 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]] @@ -85,25 +96,57 @@ def __init__(self, diffusion_config): betas_t = 1 - cumprod_alphas_t / cumprod_alphas_t_1 Q_t_list = [] - for t_idx in range(self.num_steps): - Q_t_list.append(self._uniform_transition_mat(self.num_types, betas_t[t_idx])) + 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, self.num_steps): - cumprod_Q_t_list.append(paddle.matmul(cumprod_Q_t_list[-1], Q_t[t_idx])) + 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(Q_t.shape[1], dtype="float32").unsqueeze(0), Q_t[:-1]], axis=0, + [paddle.eye(num_types, dtype="float32").unsqueeze(0), Q_t[:-1]], + axis=0, ) - self.Q_t = Q_t.reshape([-1, self.num_types, self.num_types]) - self.Q_t_1 = Q_t_1.reshape([-1, self.num_types, self.num_types]) + cumprod_Q_t_1 = paddle.concat( - [paddle.eye(cumprod_Q_t.shape[1], dtype="float32").unsqueeze(0), cumprod_Q_t[:-1]], axis=0, + [ + 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.cumprod_Q_t = cumprod_Q_t.reshape([-1, self.num_types, self.num_types]) - self.cumprod_Q_t_1 = cumprod_Q_t_1.reshape([-1, self.num_types, self.num_types]) + 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) diff --git a/structure_generation/configs/miad/README.md b/structure_generation/configs/miad/README.md index 359b0db0..a0360880 100644 --- a/structure_generation/configs/miad/README.md +++ b/structure_generation/configs/miad/README.md @@ -6,8 +6,6 @@ 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. -![MiAD Overview](https://github.com/andrey-okhotin/miad/blob/main/pictures_from_paper/miad_method_scheme.png) - --- ## Model Description @@ -24,6 +22,7 @@ MiAD defines separate forward corruption processes for $(L, F, A)$ and trains a ### Method #### 1) Lattice diffusion + The lattice is diffused with a standard DDPM Gaussian process. Supports Flow Matching as an alternative: $$ @@ -31,6 +30,7 @@ L_t = \sqrt{\bar{\alpha}_t}\,L_0 + \sqrt{1 - \bar{\alpha}_t}\,\epsilon,\quad \ep $$ #### 2) Fractional-coordinate diffusion on a torus + Fractional coordinates live on a 3D torus $[0,1)^3$. Wrapped Normal noise is used: $$ @@ -40,19 +40,24 @@ $$ 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 +## Dataset Description -| Source | Link | -|-----------------------|------| -| Original Google Drive | [Google Drive](https://drive.google.com/file/d/1BLI3VtvzfIIXlH6UHQ4o-gQaCIOZ1UR7/view?usp=sharing) | -| Mirror AiStudio | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) | +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`. @@ -60,47 +65,56 @@ Extract to `./data/mp_20/` so that CSV files are at `./data/mp_20/train.csv`, `. ## Results -| Metric | CHGNet | eq-V2 | -|--------|--------|-------| -| S.U.N. | 12.21% | 5.64% | +| 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) | --- -## Commands +## Command ### Training ```bash -# single GPU +# single-gpu training python structure_generation/train.py -c structure_generation/configs/miad/miad_mp20.yaml - -# multi-GPU (example: 4 GPUs) -python -m paddle.distributed.launch --gpus="0,1,2,3" 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' ``` -### Sampling +### Testing ```bash -# Mode 1: pre-trained model (auto-download) +# 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 2: custom checkpoint -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' -``` +# 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' -### Checkpoints +# 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 -| Source | Link | -|--------------------------------|------| -| Original (PyTorch format) | [Google Drive](https://drive.google.com/file/d/1KyD6KzvjYFPfU8lutFyO_0b8EbeHSGqf/view?usp=sharing) | -| Mirror (PyTorch format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48578/intro) | -| PaddleMaterials(Paddle format) | [AIStudio](https://aistudio.baidu.com/modelsdetail/48638/intro) | +# 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' +``` --- diff --git a/structure_generation/configs/miad/miad_mp20.yaml b/structure_generation/configs/miad/miad_mp20.yaml index e9fca929..8446a00f 100644 --- a/structure_generation/configs/miad/miad_mp20.yaml +++ b/structure_generation/configs/miad/miad_mp20.yaml @@ -50,13 +50,22 @@ Model: num_steps: 1000 lat_diffusion: method: ddpm - scheduler: diffcsp_cosine + scheduler_cfg: + __class_name__: DDPMScheduler + __init_params__: + num_train_timesteps: 1000 + beta_schedule: squaredcos_cap_v2 frac_diffusion: method: wrapped_normal - scheduler: default_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 - scheduler: default_d3pm Optimizer: __class_name__: Adam 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 index 454dfebe..2312e29c 100644 --- a/test/miad/test_miad.py +++ b/test/miad/test_miad.py @@ -44,9 +44,29 @@ "cont_time": False, "num_steps": 10, "time_embed_dim": 32, - "lat_diffusion": {"method": "ddpm", "scheduler": "diffcsp_cosine"}, - "frac_diffusion": {"method": "wrapped_normal", "scheduler": "default_wrapped_normal"}, - "type_diffusion": {"method": "d3pm", "scheduler": "default_d3pm"}, + "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"}, }