diff --git a/README.md b/README.md index ddcc5c78..ae2354d6 100755 --- a/README.md +++ b/README.md @@ -6,176 +6,125 @@ ## 🚀 Introduction -**PaddleMaterials** is an end-to-end AI4Materials toolkit built on the **PaddlePaddle** deep learning framework. Designed as a data-mechanism dual-driven platform for developing and deploying foundation models in materials science, **PPMat** enables researchers to efficiently build AI models and accelerate material discovery using pretrained models. +**PaddleMaterials** is a data-mechanism dual-driven, development and deployment of AI4Materials foundation models, end to end toolkit based on PaddlePaddle deep learning framework for materials science and engineering. **PPMat** (represents PaddleMaterials in the following text) is designed to help researchers more efficiently build AI4Materials foundation models and explore, discover, and develop new materials based on deployed pretrained models. **PPMat** has supported inorganic materials and part of organic molecules, and will support more types of materials including polymers, organic molecules, catalysts, and so on. It has supported some representative models including the equivalent graph networks-based model, diffusion model, multi-modal model, and will support more kinds of deep learing models and agents works related to AI4Material fields in the feature.

-### Core Capabilities +**Inorganic materials**, characterized by their symmetrical and periodic structures, exhibit a wide range of properties and are widely applied in various fields, from electronic devices to energy applications. Traditional experimental and computational methods for discovering crystalline materials are often time-consuming and expensive. Data-driven approaches to material discovery have the power to model the highly complex atomic systems within crystalline materials, paving the way for rapid and accurate material discovery. -| 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 | -| **Spectrum Elucidation (SE)** | Reconstruct structures from spectra | NMR structure elucidation | +**Organic materials**, distinguished by covalently linked, directionally bonded networks, mainly defined as a carbon–hydrogen or carbon–carbon bond chemical compound. These traits support core applications including flexible displays, organic photovoltaics, high-energy-density battery electrodes, advanced separation membranes, catalyts. The vast compositional and conformational space of organic molecules makes trial-and-error synthesis and ab-initio simulations slow and costly. Data-driven methods that fuse high-throughput datasets, graph-based representations, and deep generative models rapidly learn structure–property links, enabling fast virtual screening and rational design for more agile, sustainable advances in organic materials. -### Supported Materials +**Polymer materials**, characterized by their large molecular weight and complex molecular structures and built from long-chain macromolecules with tunable architectures (homopolymer, block, graft) and morphologies (amorphous, semicrystalline, cross-linked), offer lightweight, processable, and programmable mechanical, thermal, optical, and transport properties for coatings, membranes, composites, and flexible electronics. The combinatorial design space—monomer choice, sequence, tacticity, molecular-weight distribution, additives, and processing history—plus multi-scale physics makes Edisonian discovery and brute-force simulation slow and costly. Data-driven polymer informatics that fuses high-throughput measurements with graph/sequence representations and physics-guided neural surrogates learns structure-processing-property links, while generative and active-learning workflows target Tg, modulus, permeability, dielectric constant, and recyclability for rapid, sustainable polymer discovery. -- **Inorganic Crystals** - Well-supported with multiple datasets (MP2018, MP2024, JARVIS) and pretrained models -- **Organic Molecules** - Support for small molecule datasets (QM9) and property prediction -- *Polymers, catalysts, and amorphous materials are under development* +**Catalysts materials**, as key components in chemical reactions, play a crucial role in the development of new materials and technologies and spanning heterogeneous surfaces (metals, alloys, oxides, zeolites), homogeneous/organometallic complexes, and electrocatalysts, control reaction rates and selectivity across chemicals, energy, and environmental remediation. Discovery is hampered by vast compositional/structural spaces, site heterogeneity, competing pathways, and operando effects (adsorption, kinetics, deactivation) that challenge trial-and-error and exhaustive DFT. Data-driven methods—surrogate models for adsorption energies and barriers (e.g., graph neural networks), learned electronic/structural descriptors, and generative design coupled with Bayesian/active learning and automated experimentation—enable fast screening and rational optimization, accelerating catalysts for CO₂ reduction, ammonia synthesis, fuel-cell reactions, and selective oxidations. -### Why PaddleMaterials? - -- ✅ **Rich Pretrained Models** - 50+ pretrained models ready for inference -- ✅ **Multi-Task Integration** - Unified framework across PP, SG, MLIP, MLES, SE -- ✅ **Domestic Hardware Support** - Full support for MetaX GPUs and NVIDIA GPUs -- ✅ **PaddlePaddle Ecosystem** - Seamless integration with PaddlePaddle tools -- ✅ **Production-Ready** - Distributed training, mixed precision, checkpoint recovery - ---- +**Amorphous materials**, have no detectable crystal structure.its characteristic of atomic arrangement is more like liguid and has no long-range periodicity. It has attracted increasing attention duo to its broad applciations in optoelectronics, catalysis, and batteries. Its structure-property relationship is highly complex and sensitive to disorder, making it challenging to predict and design. ## 📣 News +🔥 **2025.09.25**: The **MetaX** has supported all models of multiple tasks including MLIP, MLES, PP, SG, SE. Welcome to run PaddleMaterials on MetaX chips. To experience the MetaX chip in a public cloud environment, please refer to this [PaddleMaterials_MetaX_README](./docs/MetaX/PaddleMaterials_MetaX_README.md). Pleare reference to [SupportedHardwareList](./docs/multi_device.md) for more multi-hardware adaption information. ---- +🔥 **2025.09.12**: The **Suzhou Laboratory** has established a novel model DiffNMR based on PaddleMaterials, a novel end-to-end framework that leverages a conditional discrete diffusion model for de novo molecular structure elucidation from NMR spectra. For more information, please refer to [DiffNMR](./research/DiffNMR/README.md). -## 📑 Tasks +🔥 **2025.07.01**: The **Suzhou Laboratory** has established a novel framework based on PaddleMaterials, combining an active learning workflow with conditional-diffusion-based structure generation, thereby achieving unprecedented expansion of two-dimensional material databases. For more information, please refer to [ML2DDB](./research/ML2DDB/README.md). -| Task | Description | Link | -|------|-------------|------| -| **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) | -| **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) | - ---- +## 📑 Task +- [MLIP-Machine Learning Interatomic Potential](interatomic_potentials/README.md) +- [MLES-Machine Learning Electronic Structure](electronic_structure/README.md) +- [PP-Property Prediction](property_prediction/README.md) +- [SG-Structure Generation](structure_generation/README.md) +- [SE-Spectrum Elucidation](spectrum_elucidation/README.md) ## 🔧 Installation -Please refer to the installation [document](Install.md) for your hardware environment. See [SupportedHardwareList](./docs/multi_device.md) for more multi-hardware adaptation information. +Please refer to the installation [document](Install.md) on your harware environment reference to [SupportedHardwareList](./docs/multi_device.md). ---- ## ⚡ Get Started -### Property Prediction - -Predict material formation energy using a pretrained MEGNet model: - -```bash -python property_prediction/predict.py \ - --model_name='megnet_mp2018_train_60k_e_form' \ - --weights_name='best.pdparams' \ - --cif_file_path='./property_prediction/example_data/cifs/' \ - --save_path='result.csv' -``` - -### Structure Generation - -Generate novel crystal structures: - -```bash -python structure_generation/predict.py \ - --model_name='mattergen_mp20' \ - --num_structures=100 \ - --save_path='generated_structures/' -``` - -### Interatomic Potentials - -Run molecular dynamics with ML potentials: - +PaddleMaterials offers multiple built-in models that can be directly used for inference. Taking the `megnet_mp2018_train_60k_e_form` model as an example (a MEGNet model trained on the MP2018 dataset for material formation energy prediction), use the following command for inference: ```bash -python interatomic_potentials/run_md.py - --model_name='mattersim_1M' - --structure_path='input.cif' - --temperature=300 +python property_prediction/predict.py --model_name='megnet_mp2018_train_60k_e_form' --weights_name='best.pdparams' --cif_file_path='./property_prediction/example_data/cifs/' --save_path='result.csv' ``` ---- - -### Train Your Own Model - -For training and fine-tuning, refer to the [documentation](get_started.md). - -### Contribute to PaddleMaterials - -For developer, please refer to [architecture](docs/ARCHITECTURE_ch.md). - ---- - -## 🎯 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 | - -Full model list: See [MODEL_REGISTRY](ppmat/models/__init__.py) - ---- - -## ⭐️ Star History - -[![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) - ---- + + + + + + + + + + + + + + + + + + + + + + + + + +
ParameterDescription
--model_nameName of the built-in model
--weights_nameWeights file name
--cif_file_pathPath to CIF files for prediction
--save_pathPath to save prediction results
+ +For more information on how to use PaddleMaterials to train and fine tune a model, please refer to the [documentation](get_started.md). ## 👩‍👩‍👧‍👦 Cooperation

- - - + + +

---- - ## 👩‍👩‍👧‍👦 Community -Join the PaddleMaterials WeChat group to discuss with us! +Join PaddleMaterials WeChat group to disscuss with us!

---- +## 🔄 Feedback + +We sincerely invite you to spare a moment from your busy schedule to share your [feedback](https://paddle.wjx.cn/vm/rXyQwB2.aspx#). ## 📜 License PaddleMaterials is licensed under the [Apache License 2.0](LICENSE). ---- ## 🎓 Citation -```bibtex -@misc{paddlematerials2025, - title={PaddleMaterials, a deep learning toolkit based on PaddlePaddle for material science.}, - author={PaddleMaterials Contributors}, - howpublished = {\url{https://github.com/PaddlePaddle/PaddleMaterials}}, - year={2025} -} -``` ---- + @misc{paddlematerials2025, + title={PaddleMaterials, a deep learning toolkit based on PaddlePaddle for material science.}, + author={PaddleMaterials Contributors}, + howpublished = {\url{https://github.com/PaddlePaddle/PaddleMaterials}}, + year={2025} + } -## Acknowledgements -This repository references code from the following projects: +## Acknowledgements -[PaddleScience](https://github.com/PaddlePaddle/PaddleScience) | -[Matgl](https://github.com/materialsvirtuallab/matgl) | -[CDVAE](https://github.com/txie-93/cdvae) | -[DiffCSP](https://github.com/jiaor17/DiffCSP) | -[MatterGen](https://github.com/microsoft/mattergen) | -[MatterSim](https://github.com/microsoft/mattersim) | -[CHGNet](https://github.com/CederGroupHub/chgnet) | -[AIRS](https://github.com/divelab/AIRS) \ No newline at end of file +This repository references the code from the following repositories: +[PaddleScience](https://github.com/PaddlePaddle/PaddleScience), +[Matgl](https://github.com/materialsvirtuallab/matgl), +[CDVAE](https://github.com/txie-93/cdvae), +[DiffCSP](https://github.com/jiaor17/DiffCSP), +[MatterGen](https://github.com/microsoft/mattergen), +[MatterSim](https://github.com/microsoft/mattersim), +[CHGNet](https://github.com/CederGroupHub/chgnet), +[AIRS](https://github.com/divelab/AIRS), +etc. + + diff --git a/docs/MetaX.png b/docs/MetaX.png new file mode 100644 index 00000000..585a224e Binary files /dev/null and b/docs/MetaX.png differ diff --git a/docs/suzhoulab.png b/docs/suzhoulab.png new file mode 100644 index 00000000..13e72e33 Binary files /dev/null and b/docs/suzhoulab.png differ diff --git a/docs/zhonghua.jpeg b/docs/zhonghua.jpeg new file mode 100644 index 00000000..015845f0 Binary files /dev/null and b/docs/zhonghua.jpeg differ diff --git a/interatomic_potentials/configs/schnet/schnet_md17_ethanol.yaml b/interatomic_potentials/configs/schnet/schnet_md17_ethanol.yaml new file mode 100644 index 00000000..3fa0d418 --- /dev/null +++ b/interatomic_potentials/configs/schnet/schnet_md17_ethanol.yaml @@ -0,0 +1,107 @@ +Global: + do_train: True + do_eval: True + do_test: False + + label_names: ['energy'] + + graph_converter: + __class_name__: FindPointsInSpheres + __init_params__: + cutoff: 5.0 + + prim_eager_enabled: True + + +Trainer: + max_epochs: 500 + seed: 42 + output_dir: ./output/schnet_md17_ethanol + save_freq: 50 + log_freq: 50 + + start_eval_epoch: 1 + eval_freq: 5 + pretrained_model_path: null + pretrained_weight_name: null + resume_from_checkpoint: null + use_amp: False + eval_with_no_grad: True + gradient_accumulation_steps: 1 + + best_metric_indicator: 'eval_metric' + name_for_best_metric: "energy" + greater_is_better: False + + +Model: + __class_name__: SchNet + __init_params__: + n_atom_basis: 64 + n_interactions: 6 + n_filters: 64 + cutoff: 5.0 + n_gaussians: 25 + max_z: 100 + readout: "sum" + property_names: ${Global.label_names} + data_mean: 0.0 + data_std: 1.0 + loss_type: "l1_loss" + compute_forces: False + + +Optimizer: + __class_name__: Adam + __init_params__: + lr: + __class_name__: Cosine + __init_params__: + learning_rate: 1e-4 + eta_min: 1e-7 + by_epoch: False + + +Metric: + energy: + __class_name__: IgnoreNanMetricWrapper + __init_params__: + __class_name__: paddle.nn.L1Loss + __init_params__: {} + + +Dataset: + train: + dataset: + __class_name__: MD17Dataset + __init_params__: + path: "./data/md17" + molecule: "ethanol" + property_names: ${Global.label_names} + build_graph_cfg: ${Global.graph_converter} + max_samples: 50000 + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 64 + val: + dataset: + __class_name__: MD17Dataset + __init_params__: + path: "./data/md17" + molecule: "ethanol" + property_names: ${Global.label_names} + build_graph_cfg: ${Global.graph_converter} + max_samples: 10000 + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 64 diff --git a/interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml b/interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml new file mode 100644 index 00000000..88b50ad3 --- /dev/null +++ b/interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml @@ -0,0 +1,109 @@ +Global: + do_train: True + do_eval: True + do_test: False + + label_names: ['energy_U0'] + + graph_converter: + __class_name__: FindPointsInSpheres + __init_params__: + cutoff: 10.0 + + prim_eager_enabled: True + + +Trainer: + max_epochs: 200 + seed: 42 + output_dir: ./output/schnet_qm9_U0 + save_freq: 20 + log_freq: 50 + + start_eval_epoch: 1 + eval_freq: 5 + pretrained_model_path: null + pretrained_weight_name: null + resume_from_checkpoint: null + use_amp: False + eval_with_no_grad: True + gradient_accumulation_steps: 1 + + best_metric_indicator: 'eval_metric' + name_for_best_metric: "energy_U0" + greater_is_better: False + + +Model: + __class_name__: SchNet + __init_params__: + n_atom_basis: 128 + n_interactions: 6 + n_filters: 128 + cutoff: 10.0 + n_gaussians: 50 + max_z: 100 + readout: "sum" + property_names: ${Global.label_names} + data_mean: -76.1160 + data_std: 10.3238 + loss_type: "l1_loss" + compute_forces: False + + +Optimizer: + __class_name__: Adam + __init_params__: + lr: + __class_name__: Cosine + __init_params__: + learning_rate: 1e-4 + eta_min: 1e-7 + by_epoch: False + + +Metric: + energy_U0: + __class_name__: IgnoreNanMetricWrapper + __init_params__: + __class_name__: paddle.nn.L1Loss + __init_params__: {} + + +Dataset: + train: + dataset: + __class_name__: QM9Dataset + __init_params__: + path: "./data/qm9" + property_names: ${Global.label_names} + build_graph_cfg: ${Global.graph_converter} + cache_path: "./data/qm9" + overwrite: False + filter_unvalid: True + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 64 + val: + dataset: + __class_name__: QM9Dataset + __init_params__: + path: "./data/qm9" + property_names: ${Global.label_names} + build_graph_cfg: ${Global.graph_converter} + cache_path: "./data/qm9" + overwrite: False + filter_unvalid: True + num_workers: 4 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 64 diff --git a/ppmat/datasets/alloy_dataset.py b/ppmat/datasets/alloy_dataset.py new file mode 100644 index 00000000..545e609b --- /dev/null +++ b/ppmat/datasets/alloy_dataset.py @@ -0,0 +1,88 @@ +# 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. +""" +AlloyDataset — tabular dataset for metallic glass alloy compositions. + +Loads Alloy_train.csv produced by tools/prepare_alloy_data.py. +Each sample is a 66-dimensional float vector: + columns 0-39: element composition fractions (40 elements) + columns 40-42: Tg, Tx, Tl (thermal transition temperatures in K) + columns 43-65: 23 GFA criteria (derived from Tg/Tx/Tl) + +The "source" column is dropped on load (same as original AlloyGAN). +""" + +import numpy as np +import paddle +from paddle.io import Dataset + +from ppmat.utils import logger + + +class AlloyDataset(Dataset): + """Tabular dataset for AlloyGAN training. + + Args: + path: Path to Alloy_train.csv. + categories: Optional list of dominant-element categories to filter + (e.g., ["Cu", "Fe", "Ti", "Zr"]). Default uses all entries. + normalize: Whether to normalize composition fractions to [0, 1]. + Default True (divides compositions by 100). + """ + + # Top 40 elements in order (matches CSV columns 0-39) + ELEMENTS = [ + "Cu", "Zr", "Al", "Ni", "Ti", "Ag", "Fe", "Mg", "B", "Si", + "Nb", "Y", "Ca", "La", "Co", "Be", "C", "Mo", "Pd", "P", + "Sn", "Cr", "Hf", "Zn", "Gd", "Ce", "Er", "Ga", "Au", "Nd", + "Dy", "W", "Pr", "Ta", "Sc", "Li", "Sm", "S", "Pt", "Mn", + ] + + def __init__(self, path, categories=None, normalize=True): + super().__init__() + import pandas as pd + + df = pd.read_csv(path) + + # Drop the "source" column if present (same as original code) + if "source" in df.columns: + df = df.drop(columns=["source"]) + + # Optional category filtering by dominant element + if categories is not None: + elem_cols = df.columns[:40] + dominant = df[elem_cols].idxmax(axis=1) + mask = dominant.isin(categories) + df = df[mask].reset_index(drop=True) + logger.info( + f"Filtered to categories {categories}: " + f"{len(df)} entries" + ) + + self.data = df.values.astype(np.float32) + + if normalize: + # Normalize composition fractions (0-100) to (0-1) + self.data[:, :40] = self.data[:, :40] / 100.0 + + logger.info( + f"Loaded AlloyDataset: {len(self.data)} samples, " + f"{self.data.shape[1]} features from {path}" + ) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return {"data": self.data[idx]} diff --git a/ppmat/datasets/md17_dataset.py b/ppmat/datasets/md17_dataset.py new file mode 100644 index 00000000..d75a632a --- /dev/null +++ b/ppmat/datasets/md17_dataset.py @@ -0,0 +1,250 @@ +# 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. + +"""MD17 dataset for molecular dynamics trajectories. + +The MD17 dataset (Chmiela et al., 2017) contains ab-initio molecular dynamics +trajectories for small organic molecules. Each snapshot includes atomic +positions, total energy, and per-atom forces. + +Reference: + S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky, K. T. Schütt, + K.-R. Müller. Machine Learning of Accurate Energy-Conserving Molecular + Force Fields. Science Advances, 2017. +""" + +from __future__ import annotations + +import os +import os.path as osp +import pickle +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import paddle.distributed as dist +from paddle.io import Dataset + +from ppmat.models import build_graph_converter +from ppmat.utils import download, logger + +try: + from pymatgen.core import Lattice, Structure +except ImportError: + Structure = None + Lattice = None + + +# Mapping from molecule name to original MD17 NPZ filename. +MD17_FILES = { + "benzene": "md17_benzene2017.npz", + "uracil": "md17_uracil.npz", + "naphthalene": "md17_naphthalene.npz", + "aspirin": "md17_aspirin.npz", + "salicylic_acid": "md17_salicylic.npz", + "malonaldehyde": "md17_malonaldehyde.npz", + "ethanol": "md17_ethanol.npz", + "toluene": "md17_toluene.npz", +} + +# Atomic number → element symbol (for pymatgen Structure creation). +_Z_TO_SYMBOL = { + 1: "H", 6: "C", 7: "N", 8: "O", 9: "F", 16: "S", +} + + +class MD17Dataset(Dataset): + """MD17 molecular dynamics trajectory dataset. + + Loads an MD17 NPZ file and converts each snapshot into a dict compatible + with PaddleMaterials models (optionally building a PGL graph via the + graph converter). + + The NPZ files contain: + - ``z``: atomic numbers, shape [num_atoms] + - ``R``: positions, shape [num_snapshots, num_atoms, 3] (Angstrom) + - ``E``: energies, shape [num_snapshots, 1] (kcal/mol) + - ``F``: forces, shape [num_snapshots, num_atoms, 3] (kcal/mol/A) + + Args: + path (str): Root directory for dataset storage. + molecule (str): Molecule name (e.g. "ethanol"). + property_names (str or list): Target property name(s). Default: "energy". + build_graph_cfg (dict, optional): Config for graph converter. + max_samples (int, optional): Limit dataset size (for faster debugging). + url (str, optional): Custom download URL. Default: BCS mirror. + box_size (float): Side length (A) of the cubic cell used for + non-periodic molecules. Default: 100.0. + cache_graphs (bool): Whether to cache built graphs to disk. Default: True. + """ + + # BCS mirror for MD17 NPZ files + default_url = "https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MD17" + + def __init__( + self, + path: str, + molecule: str = "ethanol", + property_names: Union[str, List[str]] = "energy", + *, + build_graph_cfg: Optional[Dict] = None, + max_samples: Optional[int] = None, + url: Optional[str] = None, + box_size: float = 100.0, + cache_graphs: bool = True, + **kwargs, + ) -> None: + super().__init__() + + if molecule not in MD17_FILES: + raise ValueError( + f"Unknown molecule '{molecule}'. Choose from: {list(MD17_FILES)}" + ) + + if isinstance(property_names, str): + property_names = [property_names] + self.property_names = property_names + self.molecule = molecule + self.box_size = box_size + self.build_graph_cfg = build_graph_cfg + + # Paths + os.makedirs(path, exist_ok=True) + self.raw_dir = osp.join(path, "raw_md17") + os.makedirs(self.raw_dir, exist_ok=True) + + npz_name = MD17_FILES[molecule] + npz_path = osp.join(self.raw_dir, npz_name) + + # Download if needed + if not osp.exists(npz_path): + base_url = url or self.default_url + full_url = f"{base_url}/{npz_name}" + logger.info(f"Downloading MD17 {molecule} from {full_url}") + download.download_file(full_url, npz_path) + + # Load raw data + raw = np.load(npz_path) + self.atomic_numbers = raw["z"].astype(np.int64) # [num_atoms] + self.positions = raw["R"].astype(np.float32) # [N, num_atoms, 3] + self.energies = raw["E"].astype(np.float32) # [N, 1] or [N] + self.forces = raw["F"].astype(np.float32) # [N, num_atoms, 3] + + if self.energies.ndim == 1: + self.energies = self.energies[:, None] + + if max_samples is not None: + self.positions = self.positions[:max_samples] + self.energies = self.energies[:max_samples] + self.forces = self.forces[:max_samples] + + self.num_samples = len(self.positions) + logger.info( + f"MD17 {molecule}: {self.num_samples} snapshots, " + f"{len(self.atomic_numbers)} atoms/snapshot" + ) + + # Build and cache graphs if configured + self.graphs = None + if build_graph_cfg is not None: + graph_converter_name = build_graph_cfg.get("__class_name__", "custom") + cutoff = build_graph_cfg.get("__init_params__", {}).get("cutoff", 5) + cache_dir = osp.join( + path, + f"md17_{molecule}_cache_{graph_converter_name}_cutoff_{int(cutoff)}", + "graphs", + ) + + done_flag = osp.join(cache_dir, "completed.flag") + if cache_graphs and osp.exists(done_flag): + logger.info(f"Loading cached graphs from {cache_dir}") + self.graphs = self._load_cached_graphs(cache_dir) + else: + if dist.get_rank() == 0: + logger.info("Building graphs for MD17 dataset...") + os.makedirs(cache_dir, exist_ok=True) + converter = build_graph_converter(build_graph_cfg) + structures = self._build_structures() + self.graphs = converter(structures) + if cache_graphs: + self._save_cached_graphs(cache_dir) + with open(done_flag, "w") as f: + f.write("done") + if dist.is_initialized(): + dist.barrier() + if self.graphs is None: + self.graphs = self._load_cached_graphs(cache_dir) + + def _build_structures(self) -> list: + """Convert all snapshots to pymatgen Structures.""" + if Structure is None: + raise RuntimeError("pymatgen is required: pip install pymatgen") + + lattice = Lattice.from_parameters( + self.box_size, self.box_size, self.box_size, 90, 90, 90 + ) + species = [_Z_TO_SYMBOL.get(z, str(z)) for z in self.atomic_numbers] + + structures = [] + for i in range(self.num_samples): + coords = self.positions[i] # [num_atoms, 3] + struct = Structure( + lattice, + species, + coords, + coords_are_cartesian=True, + ) + structures.append(struct) + return structures + + def _save_cached_graphs(self, cache_dir: str) -> None: + for i, g in enumerate(self.graphs): + with open(osp.join(cache_dir, f"{i:08d}.pkl"), "wb") as f: + pickle.dump(g, f) + + def _load_cached_graphs(self, cache_dir: str) -> list: + files = sorted( + f for f in os.listdir(cache_dir) if f.endswith(".pkl") + ) + graphs = [] + for fn in files: + with open(osp.join(cache_dir, fn), "rb") as f: + graphs.append(pickle.load(f)) + return graphs + + def __len__(self) -> int: + return self.num_samples + + def __getitem__(self, idx: int) -> Dict[str, Any]: + data = {} + + if self.graphs is not None: + data["graph"] = self.graphs[idx] + else: + # Return raw data without graph (fallback) + data["pos"] = self.positions[idx] + data["atomic_numbers"] = self.atomic_numbers.copy() + data["cell"] = np.eye(3, dtype="float32") * self.box_size + data["natoms"] = len(self.atomic_numbers) + data["pbc"] = np.array([False, False, False], dtype=bool) + + # Properties + for pname in self.property_names: + if pname == "energy": + data["energy"] = self.energies[idx] + elif pname == "forces": + data["forces"] = self.forces[idx] + else: + data[pname] = self.energies[idx] + + return data diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py index 95d73232..fd292087 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.diffsyn.diffsyn import DiffSyn from ppmat.utils import download from ppmat.utils import logger from ppmat.utils import save_load @@ -67,6 +68,7 @@ "DiffNMR", "InfGCN", "MatENO", + "DiffSyn", ] # Warning: The key of the dictionary must be consistent with the file name of the value diff --git a/ppmat/models/diffsyn/__init__.py b/ppmat/models/diffsyn/__init__.py new file mode 100644 index 00000000..1947662a --- /dev/null +++ b/ppmat/models/diffsyn/__init__.py @@ -0,0 +1 @@ +from .diffsyn import DiffSyn diff --git a/ppmat/models/diffsyn/diffsyn.py b/ppmat/models/diffsyn/diffsyn.py new file mode 100644 index 00000000..51b31f32 --- /dev/null +++ b/ppmat/models/diffsyn/diffsyn.py @@ -0,0 +1,796 @@ +"""DiffSyn: Conditional 1D Diffusion for Zeolite Synthesis Conditions. + +Adapted from https://github.com/eltonpan/zeosyn_gen (PyTorch) to PaddlePaddle. +Uses a 1D U-Net denoiser wrapped in a Gaussian diffusion process (DDPM/DDIM) +to generate synthesis conditions conditioned on zeolite/OSDA features. +""" + +import math +from collections import namedtuple +from functools import partial + +import numpy as np +import paddle +import paddle.nn as nn +import paddle.nn.functional as F + +ModelPrediction = namedtuple("ModelPrediction", ["pred_noise", "pred_x_start"]) + + +# --------------------------------------------------------------------------- +# Helper utilities +# --------------------------------------------------------------------------- + +def exists(x): + return x is not None + + +def default(val, d): + if exists(val): + return val + return d() if callable(d) else d + + +def identity(t, *args, **kwargs): + return t + + +def cosine_beta_schedule(timesteps, s=0.008): + """Cosine schedule as in https://openreview.net/forum?id=-NEXDKk8gZ.""" + steps = timesteps + 1 + x = np.linspace(0, timesteps, steps, dtype=np.float64) + alphas_cumprod = np.cos((x / timesteps + s) / (1 + s) * math.pi * 0.5) ** 2 + alphas_cumprod = alphas_cumprod / alphas_cumprod[0] + betas = 1.0 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) + return np.clip(betas, 0, 0.9999) + + +def linear_beta_schedule(timesteps): + scale = 1000 / timesteps + beta_start = scale * 0.0001 + beta_end = scale * 0.02 + return np.linspace(beta_start, beta_end, timesteps, dtype=np.float64) + + +def extract(a, t, x_shape): + """Gather values from *a* at indices *t*, then broadcast to *x_shape*.""" + b = t.shape[0] + out = paddle.gather(a, t) + return out.reshape([b] + [1] * (len(x_shape) - 1)) + + +# --------------------------------------------------------------------------- +# Tiny building blocks +# --------------------------------------------------------------------------- + +class SinusoidalPosEmb(nn.Layer): + """Sinusoidal positional embedding for diffusion timesteps.""" + + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = paddle.exp(paddle.arange(half_dim, dtype="float32") * -emb) + emb = x.unsqueeze(-1).astype("float32") * emb.unsqueeze(0) + emb = paddle.concat([emb.sin(), emb.cos()], axis=-1) + return emb + + +class LayerNorm1D(nn.Layer): + """Channel-wise layer-norm for [B, C, L] tensors.""" + + def __init__(self, dim): + super().__init__() + self.g = paddle.create_parameter( + shape=[1, dim, 1], + dtype="float32", + default_initializer=nn.initializer.Constant(1.0), + ) + + def forward(self, x): + eps = 1e-5 if x.dtype == paddle.float32 else 1e-3 + var = x.var(axis=1, keepdim=True) + mean = x.mean(axis=1, keepdim=True) + return (x - mean) * (var + eps).rsqrt() * self.g + + +class Residual(nn.Layer): + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, x, *args, **kwargs): + return self.fn(x, *args, **kwargs) + x + + +class PreNorm(nn.Layer): + def __init__(self, dim, fn): + super().__init__() + self.fn = fn + self.norm = LayerNorm1D(dim) + + def forward(self, x): + x = self.norm(x) + return self.fn(x) + + +# --------------------------------------------------------------------------- +# Weight-standardised Conv1D +# --------------------------------------------------------------------------- + +class WSConv1D(nn.Layer): + """Conv1D with weight standardisation (synergistic with GroupNorm).""" + + def __init__(self, in_channels, out_channels, kernel_size, + stride=1, padding=0, groups=1, bias_attr=True): + super().__init__() + self.conv = nn.Conv1D( + in_channels, out_channels, kernel_size, + stride=stride, padding=padding, groups=groups, + bias_attr=bias_attr, + ) + + def forward(self, x): + eps = 1e-5 if x.dtype == paddle.float32 else 1e-3 + weight = self.conv.weight # [out, in/groups, k] + mean = weight.mean(axis=[1, 2], keepdim=True) + var = weight.var(axis=[1, 2], keepdim=True) + normed = (weight - mean) * (var + eps).rsqrt() + return F.conv1d(x, normed, self.conv.bias, self.conv._stride, + self.conv._padding, self.conv._dilation, self.conv._groups) + + +# --------------------------------------------------------------------------- +# Core blocks +# --------------------------------------------------------------------------- + +class Block1D(nn.Layer): + """Conv1D → GroupNorm → SiLU, with optional FiLM scale-shift.""" + + def __init__(self, dim_in, dim_out, groups=8): + super().__init__() + self.proj = WSConv1D(dim_in, dim_out, 3, padding=1) + self.norm = nn.GroupNorm(groups, dim_out) + self.act = nn.Silu() + + def forward(self, x, scale_shift=None): + x = self.proj(x) + x = self.norm(x) + if exists(scale_shift): + scale, shift = scale_shift + x = x * (scale + 1) + shift + x = self.act(x) + return x + + +class ResnetBlock1D(nn.Layer): + """1D ResNet block with optional time + condition FiLM modulation.""" + + def __init__(self, dim_in, dim_out, *, time_emb_dim=None, + cond_emb_dim=None, groups=8): + super().__init__() + mlp_dim = int(default(time_emb_dim, 0)) + int(default(cond_emb_dim, 0)) + self.mlp = ( + nn.Sequential(nn.Silu(), nn.Linear(mlp_dim, dim_out * 2)) + if mlp_dim > 0 + else None + ) + self.block1 = Block1D(dim_in, dim_out, groups=groups) + self.block2 = Block1D(dim_out, dim_out, groups=groups) + self.res_conv = ( + nn.Conv1D(dim_in, dim_out, 1) if dim_in != dim_out else nn.Identity() + ) + + def forward(self, x, time_emb=None, cond_emb=None): + scale_shift = None + if exists(self.mlp) and (exists(time_emb) or exists(cond_emb)): + parts = [e for e in (time_emb, cond_emb) if exists(e)] + cond = paddle.concat(parts, axis=-1) + cond = self.mlp(cond) + cond = cond.unsqueeze(-1) # [B, 2*dim_out, 1] + scale_shift = cond.chunk(2, axis=1) + h = self.block1(x, scale_shift=scale_shift) + h = self.block2(h) + return h + self.res_conv(x) + + +class LinearAttention1D(nn.Layer): + """Efficient O(n) linear attention for 1D sequences.""" + + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.scale = dim_head ** -0.5 + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv1D(dim, hidden_dim * 3, 1, bias_attr=False) + self.to_out = nn.Sequential( + nn.Conv1D(hidden_dim, dim, 1), + LayerNorm1D(dim), + ) + + def forward(self, x): + b, c, n = x.shape + h = self.heads + qkv = self.to_qkv(x).chunk(3, axis=1) + # reshape [B, hidden, N] → [B, heads, dim_head, N] + q, k, v = [ + t.reshape([b, h, -1, n]) for t in qkv + ] + q = F.softmax(q, axis=-2) + k = F.softmax(k, axis=-1) + q = q * self.scale + + # linear attention: context = k^T v, out = context^T q + context = paddle.matmul(k, v.transpose([0, 1, 3, 2])) # [B,h,d,d] + out = paddle.matmul(context, q) # [B,h,d,N] + out = out.reshape([b, -1, n]) + return self.to_out(out) + + +class Attention1D(nn.Layer): + """Standard dot-product attention for 1D sequences.""" + + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.scale = dim_head ** -0.5 + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv1D(dim, hidden_dim * 3, 1, bias_attr=False) + self.to_out = nn.Conv1D(hidden_dim, dim, 1) + + def forward(self, x): + b, c, n = x.shape + h = self.heads + qkv = self.to_qkv(x).chunk(3, axis=1) + q, k, v = [t.reshape([b, h, -1, n]) for t in qkv] + q = q * self.scale + sim = paddle.einsum("bhdi,bhdj->bhij", q, k) + attn = F.softmax(sim, axis=-1) + # out = einsum('b h i j, b h d j -> b h i d', attn, v) → [B,h,N,d_h] + out = paddle.einsum("bhij,bhdj->bhid", attn, v) + # reshape to [B, hidden, N] + out = out.transpose([0, 1, 3, 2]).reshape([b, -1, n]) + return self.to_out(out) + + +def _downsample_1d(dim, dim_out=None): + """Strided Conv1D downsample (factor 2).""" + return nn.Conv1D(dim, default(dim_out, dim), 4, stride=2, padding=1) + + +class _Upsample1DNearest(nn.Layer): + """Nearest-neighbour 1D upsample (factor 2) via repeat-interleave.""" + + def forward(self, x): + # x: [B, C, L] → [B, C, L*2] + return x.repeat_interleave(2, axis=2) + + +def _upsample_1d(dim, dim_out=None): + """Nearest-neighbour upsample + Conv1D.""" + return nn.Sequential( + _Upsample1DNearest(), + nn.Conv1D(dim, default(dim_out, dim), 3, padding=1), + ) + + +# --------------------------------------------------------------------------- +# 1-D U-Net +# --------------------------------------------------------------------------- + +class Unet1D(nn.Layer): + """1D U-Net denoising model with condition encoding. + + Args: + dim: base channel dimension. + channels: input / output channels (synthesis-condition dimensions). + cond_dim: raw condition feature dimension fed into the condition MLP. + dim_mults: channel multipliers per resolution level. + groups: GroupNorm groups. + cond_drop_prob: probability of dropping conditions (classifier-free guidance). + """ + + def __init__( + self, + dim=64, + channels=3, + cond_dim=32, + dim_mults=(1, 2, 4), + groups=8, + cond_drop_prob=0.5, + ): + super().__init__() + self.channels = channels + self.out_dim = channels + self.cond_drop_prob = cond_drop_prob + + # initial projection + init_dim = dim + self.init_conv = nn.Conv1D(channels, init_dim, 7, padding=3) + + dims = [init_dim] + [dim * m for m in dim_mults] + in_out = list(zip(dims[:-1], dims[1:])) + + time_dim = dim * 4 + cond_inner_dim = dim * 4 + + # time embedding + self.time_mlp = nn.Sequential( + SinusoidalPosEmb(dim), + nn.Linear(dim, time_dim), + nn.GELU(), + nn.Linear(time_dim, time_dim), + ) + + # condition MLP + self.cond_mlp = nn.Sequential( + nn.Linear(cond_dim, cond_inner_dim), + nn.ReLU(), + nn.Linear(cond_inner_dim, cond_inner_dim), + ) + + # null-condition embedding for classifier-free guidance + self.null_cond_emb = paddle.create_parameter( + shape=[cond_dim], + dtype="float32", + default_initializer=nn.initializer.Normal(mean=0.0, std=1.0), + ) + + block_klass = partial( + ResnetBlock1D, time_emb_dim=time_dim, cond_emb_dim=cond_inner_dim, groups=groups + ) + + # --- encoder (down) --- + self.downs = nn.LayerList() + for ind, (d_in, d_out) in enumerate(in_out): + is_last = ind >= len(in_out) - 1 + self.downs.append( + nn.LayerList( + [ + block_klass(d_in, d_in), + block_klass(d_in, d_in), + Residual(PreNorm(d_in, LinearAttention1D(d_in))), + _downsample_1d(d_in, d_out) + if not is_last + else nn.Conv1D(d_in, d_out, 3, padding=1), + ] + ) + ) + + mid_dim = dims[-1] + self.mid_block1 = block_klass(mid_dim, mid_dim) + self.mid_attn = Residual(PreNorm(mid_dim, Attention1D(mid_dim))) + self.mid_block2 = block_klass(mid_dim, mid_dim) + + # --- decoder (up) --- + self.ups = nn.LayerList() + for ind, (d_in, d_out) in enumerate(reversed(in_out)): + is_last = ind == len(in_out) - 1 + self.ups.append( + nn.LayerList( + [ + block_klass(d_out + d_in, d_out), + block_klass(d_out + d_in, d_out), + Residual(PreNorm(d_out, LinearAttention1D(d_out))), + _upsample_1d(d_out, d_in) + if not is_last + else nn.Conv1D(d_out, d_in, 3, padding=1), + ] + ) + ) + + self.final_res_block = block_klass(dim * 2, dim) + self.final_conv = nn.Conv1D(dim, self.out_dim, 1) + + # -- classifier-free guidance entry point -- + def forward_with_cond_scale(self, *args, cond_scale=1.0, **kwargs): + logits = self.forward(*args, cond_drop_prob=0.0, **kwargs) + if cond_scale == 1.0: + return logits + null_logits = self.forward(*args, cond_drop_prob=1.0, **kwargs) + return null_logits + (logits - null_logits) * cond_scale + + def forward(self, x, time, cond=None, cond_drop_prob=None): + """ + Args: + x: [B, C, L] noisy input. + time: [B] integer timestep indices. + cond: [B, cond_dim] conditioning features (or ``None``). + cond_drop_prob: override for condition dropout probability. + """ + batch = x.shape[0] + cond_drop_prob = default(cond_drop_prob, self.cond_drop_prob) + + # --- condition encoding with classifier-free dropout --- + if exists(cond): + if cond_drop_prob > 0.0: + keep_mask = ( + paddle.rand([batch]) >= cond_drop_prob + ) # True → keep + null_cond = self.null_cond_emb.unsqueeze(0).expand([batch, -1]) + cond = paddle.where( + keep_mask.unsqueeze(-1), cond, null_cond + ) + c = self.cond_mlp(cond) + else: + # unconditional — use null embedding for every sample + null_cond = self.null_cond_emb.unsqueeze(0).expand([batch, -1]) + c = self.cond_mlp(null_cond) + + # --- time embedding --- + t = self.time_mlp(time) + + # --- U-Net --- + x = self.init_conv(x) + r = x.clone() + + h = [] + for block1, block2, attn, downsample in self.downs: + x = block1(x, t, c) + h.append(x) + x = block2(x, t, c) + x = attn(x) + h.append(x) + x = downsample(x) + + x = self.mid_block1(x, t, c) + x = self.mid_attn(x) + x = self.mid_block2(x, t, c) + + for block1, block2, attn, upsample in self.ups: + x = paddle.concat([x, h.pop()], axis=1) + x = block1(x, t, c) + x = paddle.concat([x, h.pop()], axis=1) + x = block2(x, t, c) + x = attn(x) + x = upsample(x) + + x = paddle.concat([x, r], axis=1) + x = self.final_res_block(x, t, c) + return self.final_conv(x) + + +# --------------------------------------------------------------------------- +# DiffSyn — top-level model following PaddleMaterials conventions +# --------------------------------------------------------------------------- + +class DiffSyn(nn.Layer): + """DiffSyn: Conditional 1D Gaussian Diffusion for zeolite synthesis conditions. + + Wraps :class:`Unet1D` in a DDPM / DDIM diffusion process and exposes + the ``forward`` / ``predict`` API expected by PaddleMaterials. + + Args: + dim: base U-Net channel width. + channels: number of synthesis-condition channels. + seq_length: temporal length of the 1D signal. + cond_dim: condition feature dimensionality. + dim_mults: channel multipliers per resolution level. + timesteps: number of diffusion timesteps *T*. + beta_schedule: ``'cosine'`` or ``'linear'``. + objective: prediction target — ``'pred_noise'``, ``'pred_x0'``, or ``'pred_v'``. + cond_drop_prob: classifier-free guidance dropout probability. + property_names: key used in the returned prediction dict. + data_mean / data_std: normalisation statistics. + loss_type: ``'l1_loss'`` or ``'mse_loss'``. + """ + + def __init__( + self, + dim=64, + channels=3, + seq_length=8, + cond_dim=32, + dim_mults=(1, 2, 4), + groups=8, + timesteps=1000, + sampling_timesteps=None, + beta_schedule="cosine", + objective="pred_noise", + cond_drop_prob=0.5, + ddim_sampling_eta=1.0, + property_names="synthesis_conditions", + data_mean=0.0, + data_std=1.0, + loss_type="l1_loss", + ): + super().__init__() + + self.seq_length = seq_length + self.channels = channels + self.objective = objective + assert objective in { + "pred_noise", + "pred_x0", + "pred_v", + }, f"objective must be pred_noise | pred_x0 | pred_v, got {objective}" + + # denoising backbone + self.denoise_model = Unet1D( + dim=dim, + channels=channels, + cond_dim=cond_dim, + dim_mults=dim_mults, + groups=groups, + cond_drop_prob=cond_drop_prob, + ) + + # --- diffusion schedule --- + if beta_schedule == "cosine": + betas = cosine_beta_schedule(timesteps) + elif beta_schedule == "linear": + betas = linear_beta_schedule(timesteps) + else: + raise ValueError(f"Unknown beta schedule: {beta_schedule}") + + alphas = 1.0 - betas + alphas_cumprod = np.cumprod(alphas, axis=0) + alphas_cumprod_prev = np.concatenate([[1.0], alphas_cumprod[:-1]]) + + self.num_timesteps = int(betas.shape[0]) + + # sampling + self.sampling_timesteps = default(sampling_timesteps, self.num_timesteps) + self.is_ddim_sampling = self.sampling_timesteps < self.num_timesteps + self.ddim_sampling_eta = ddim_sampling_eta + + def _buf(name, arr): + self.register_buffer(name, paddle.to_tensor(arr, dtype="float32")) + + _buf("betas", betas) + _buf("alphas_cumprod", alphas_cumprod) + _buf("alphas_cumprod_prev", alphas_cumprod_prev) + + _buf("sqrt_alphas_cumprod", np.sqrt(alphas_cumprod)) + _buf("sqrt_one_minus_alphas_cumprod", np.sqrt(1.0 - alphas_cumprod)) + _buf("log_one_minus_alphas_cumprod", np.log(1.0 - alphas_cumprod)) + _buf("sqrt_recip_alphas_cumprod", np.sqrt(1.0 / alphas_cumprod)) + _buf("sqrt_recipm1_alphas_cumprod", np.sqrt(1.0 / alphas_cumprod - 1.0)) + + # posterior q(x_{t-1} | x_t, x_0) + posterior_variance = ( + betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod) + ) + _buf("posterior_variance", posterior_variance) + _buf( + "posterior_log_variance_clipped", + np.log(np.clip(posterior_variance, 1e-20, None)), + ) + _buf( + "posterior_mean_coef1", + betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod), + ) + _buf( + "posterior_mean_coef2", + (1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod), + ) + + # SNR-based loss weight + snr = alphas_cumprod / (1.0 - alphas_cumprod) + if objective == "pred_noise": + loss_weight = np.ones_like(snr) + elif objective == "pred_x0": + loss_weight = snr.copy() + elif objective == "pred_v": + loss_weight = snr / (snr + 1.0) + _buf("loss_weight", loss_weight) + + # PM conventions + if isinstance(property_names, list): + self.property_names = property_names[0] + else: + self.property_names = property_names + + if loss_type == "l1_loss": + self.loss_fn = F.l1_loss + elif loss_type == "mse_loss": + self.loss_fn = F.mse_loss + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + self.register_buffer( + "data_mean", paddle.to_tensor(np.atleast_1d(np.float32(data_mean))) + ) + self.register_buffer( + "data_std", paddle.to_tensor(np.atleast_1d(np.float32(data_std))) + ) + + # -- normalisation helpers (PM convention) -- + def normalize(self, t): + return (t - self.data_mean) / self.data_std + + def unnormalize(self, t): + return t * self.data_std + self.data_mean + + # ------------------------------------------------------------------ + # Diffusion math + # ------------------------------------------------------------------ + + def q_sample(self, x_start, t, noise=None): + """Forward diffusion q(x_t | x_0).""" + if noise is None: + noise = paddle.randn(x_start.shape) + return ( + extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + def predict_start_from_noise(self, x_t, t, noise): + return ( + extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t + - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise + ) + + def predict_noise_from_start(self, x_t, t, x0): + return ( + extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0 + ) / extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def predict_v(self, x_start, t, noise): + return ( + extract(self.sqrt_alphas_cumprod, t, x_start.shape) * noise + - extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * x_start + ) + + def predict_start_from_v(self, x_t, t, v): + return ( + extract(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t + - extract(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v + ) + + def q_posterior(self, x_start, x_t, t): + posterior_mean = ( + extract(self.posterior_mean_coef1, t, x_t.shape) * x_start + + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = extract(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = extract( + self.posterior_log_variance_clipped, t, x_t.shape + ) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def model_predictions(self, x, t, cond=None, cond_scale=3.0, clip_x_start=False): + model_output = self.denoise_model.forward_with_cond_scale( + x, t, cond=cond, cond_scale=cond_scale + ) + maybe_clip = ( + partial(paddle.clip, min=-1.0, max=1.0) if clip_x_start else identity + ) + + if self.objective == "pred_noise": + pred_noise = model_output + x_start = self.predict_start_from_noise(x, t, pred_noise) + x_start = maybe_clip(x_start) + elif self.objective == "pred_x0": + x_start = model_output + x_start = maybe_clip(x_start) + pred_noise = self.predict_noise_from_start(x, t, x_start) + elif self.objective == "pred_v": + v = model_output + x_start = self.predict_start_from_v(x, t, v) + x_start = maybe_clip(x_start) + pred_noise = self.predict_noise_from_start(x, t, x_start) + + return ModelPrediction(pred_noise, x_start) + + def p_mean_variance(self, x, t, cond=None, cond_scale=3.0, clip_denoised=True): + preds = self.model_predictions(x, t, cond=cond, cond_scale=cond_scale) + x_start = preds.pred_x_start + if clip_denoised: + x_start = paddle.clip(x_start, -1.0, 1.0) + model_mean, post_var, post_log_var = self.q_posterior( + x_start=x_start, x_t=x, t=t + ) + return model_mean, post_var, post_log_var, x_start + + # ------------------------------------------------------------------ + # Sampling + # ------------------------------------------------------------------ + + @paddle.no_grad() + def p_sample(self, x, t_int, cond=None, cond_scale=3.0, clip_denoised=True): + b = x.shape[0] + batched_times = paddle.full([b], t_int, dtype="int64") + model_mean, _, model_log_var, x_start = self.p_mean_variance( + x=x, t=batched_times, cond=cond, cond_scale=cond_scale, + clip_denoised=clip_denoised, + ) + noise = paddle.randn(x.shape) if t_int > 0 else paddle.zeros(x.shape) + pred = model_mean + (0.5 * model_log_var).exp() * noise + return pred, x_start + + @paddle.no_grad() + def p_sample_loop(self, shape, cond=None, cond_scale=3.0): + img = paddle.randn(shape) + for t in reversed(range(0, self.num_timesteps)): + img, _ = self.p_sample(img, t, cond=cond, cond_scale=cond_scale) + return img + + @paddle.no_grad() + def ddim_sample(self, shape, cond=None, cond_scale=3.0, clip_denoised=True): + total_timesteps = self.num_timesteps + sampling_timesteps = self.sampling_timesteps + eta = self.ddim_sampling_eta + + times = np.linspace(-1, total_timesteps - 1, sampling_timesteps + 1) + times = list(reversed(times.astype(int).tolist())) + time_pairs = list(zip(times[:-1], times[1:])) + + batch = shape[0] + img = paddle.randn(shape) + + for time, time_next in time_pairs: + time_cond = paddle.full([batch], time, dtype="int64") + pred_noise, x_start = self.model_predictions( + img, time_cond, cond=cond, cond_scale=cond_scale, + clip_x_start=clip_denoised, + ) + if time_next < 0: + img = x_start + continue + alpha = self.alphas_cumprod[time] + alpha_next = self.alphas_cumprod[time_next] + sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt() + c = (1 - alpha_next - sigma ** 2).sqrt() + noise = paddle.randn(img.shape) + img = x_start * alpha_next.sqrt() + c * pred_noise + sigma * noise + + return img + + @paddle.no_grad() + def sample(self, batch_size=1, cond=None, cond_scale=3.0): + """Full diffusion sampling (DDPM or DDIM).""" + shape = [batch_size, self.channels, self.seq_length] + sample_fn = self.ddim_sample if self.is_ddim_sampling else self.p_sample_loop + return sample_fn(shape, cond=cond, cond_scale=cond_scale) + + # ------------------------------------------------------------------ + # Training forward (PM convention) + # ------------------------------------------------------------------ + + def _forward(self, data): + """Core forward: compute diffusion training loss targets.""" + x_start = data["x"] # [B, C, L] + cond = data.get("cond", None) + b = x_start.shape[0] + t = paddle.randint(0, self.num_timesteps, [b]) + noise = paddle.randn(x_start.shape) + x_noisy = self.q_sample(x_start, t, noise) + pred = self.denoise_model(x_noisy, t, cond=cond) + + if self.objective == "pred_noise": + target = noise + elif self.objective == "pred_x0": + target = x_start + elif self.objective == "pred_v": + target = self.predict_v(x_start, t, noise) + + return pred, target + + def forward(self, data, return_loss=True, return_prediction=True): + assert return_loss or return_prediction, ( + "At least one of return_loss or return_prediction must be True." + ) + pred, target = self._forward(data) + + loss_dict = {} + if return_loss: + loss = self.loss_fn(input=pred, label=target) + loss_dict["loss"] = loss + + prediction = {} + if return_prediction: + prediction[self.property_names] = pred + + return {"loss_dict": loss_dict, "pred_dict": prediction} + + @paddle.no_grad() + def predict(self, data): + """Generate samples via the full reverse diffusion process.""" + cond = data.get("cond", None) + batch_size = cond.shape[0] if exists(cond) else 1 + samples = self.sample(batch_size=batch_size, cond=cond) + return {self.property_names: self.unnormalize(samples)} diff --git a/structure_generation/configs/diffsyn/README.md b/structure_generation/configs/diffsyn/README.md new file mode 100644 index 00000000..cad9604b --- /dev/null +++ b/structure_generation/configs/diffsyn/README.md @@ -0,0 +1,43 @@ +# DiffSyn + +[Generative Design of Inorganic Compounds Using a Deep Conditional Diffusion Model](https://github.com/eltonpan/zeosyn_gen) + +## Abstract + +DiffSyn uses a conditional 1D Gaussian diffusion process to generate zeolite synthesis conditions. A 1D U-Net denoiser with classifier-free guidance generates synthesis parameter sequences conditioned on zeolite/OSDA features. The model supports DDPM and DDIM sampling with cosine or linear noise schedules. + +## Model + +DiffSyn wraps a 1D U-Net in a DDPM/DDIM diffusion framework. The U-Net processes temporal sequences of synthesis parameters, using sinusoidal time embeddings, condition MLPs for classifier-free guidance, and multi-scale ResNet blocks with linear attention. During training, the model learns to denoise corrupted sequences; at inference, it progressively denoises random Gaussian noise into valid synthesis conditions. + +## Training + +```bash +# single-gpu training +python structure_generation/train.py -c structure_generation/configs/diffsyn/diffsyn_zeolite.yaml + +# multi-gpu training +python -m paddle.distributed.launch --gpus="0,1,2,3" structure_generation/train.py -c structure_generation/configs/diffsyn/diffsyn_zeolite.yaml +``` + +## Validation + +```bash +python structure_generation/train.py -c structure_generation/configs/diffsyn/diffsyn_zeolite.yaml Global.do_eval=True Global.do_train=False Global.do_test=False Trainer.pretrained_model_path='path/to/model.pdparams' +``` + +## Testing + +```bash +python structure_generation/train.py -c structure_generation/configs/diffsyn/diffsyn_zeolite.yaml Global.do_eval=False Global.do_train=False Global.do_test=True Trainer.pretrained_model_path='path/to/model.pdparams' +``` + +## Citation + +``` +@article{pan2024diffsyn, + title={Generative design of inorganic compounds using a deep conditional diffusion model}, + author={Pan, Elton and Kwon, Soonhyoung and Xie, Tian}, + year={2024} +} +``` diff --git a/structure_generation/configs/diffsyn/diffsyn_zeolite.yaml b/structure_generation/configs/diffsyn/diffsyn_zeolite.yaml new file mode 100644 index 00000000..3ec1b5c5 --- /dev/null +++ b/structure_generation/configs/diffsyn/diffsyn_zeolite.yaml @@ -0,0 +1,102 @@ +Global: + do_train: True + do_eval: False + do_test: False + num_train_timesteps: 1000 + +Trainer: + max_epochs: 1000 + seed: 42 + output_dir: ./output/diffsyn_zeolite + 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__: DiffSyn + __init_params__: + dim: 64 + channels: 3 + seq_length: 8 + cond_dim: 32 + dim_mults: [1, 2, 4] + groups: 8 + timesteps: ${Global.num_train_timesteps} + beta_schedule: cosine + objective: pred_noise + cond_drop_prob: 0.5 + ddim_sampling_eta: 1.0 + loss_type: l1_loss + +Optimizer: + __class_name__: Adam + __init_params__: + beta1: 0.9 + beta2: 0.999 + lr: + __class_name__: ReduceOnPlateau + __init_params__: + learning_rate: 0.001 + factor: 0.5 + by_epoch: True + patience: 50 + min_lr: 0.0001 + indicator: "train_loss" + indicator_name: 'loss' + +Dataset: + train: + dataset: + __class_name__: ZeoliteDataset + __init_params__: + path: "./data/zeolite/train.csv" + loader: + num_workers: 0 + use_shared_memory: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 128 + val: + dataset: + __class_name__: ZeoliteDataset + __init_params__: + path: "./data/zeolite/val.csv" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 64 + test: + dataset: + __class_name__: ZeoliteDataset + __init_params__: + path: "./data/zeolite/test.csv" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 64 + +Sample: + model_sample_params: + num_inference_steps: 1000 diff --git a/test/diffsyn/__init__.py b/test/diffsyn/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/diffsyn/conftest.py b/test/diffsyn/conftest.py new file mode 100644 index 00000000..7b473321 --- /dev/null +++ b/test/diffsyn/conftest.py @@ -0,0 +1,23 @@ +# 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. + +"""Conftest for DiffSyn tests — stubs heavy dependencies for CPU testing.""" + +import sys +import types + +# Stub pgl (PaddlePaddle Graph Learning) — not needed for unit tests +pgl_stub = types.ModuleType("pgl") +pgl_stub.Graph = type("Graph", (), {}) +sys.modules.setdefault("pgl", pgl_stub) diff --git a/test/diffsyn/raw_infer_data/__init__.py b/test/diffsyn/raw_infer_data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/diffsyn/raw_infer_data/create_input_output_npz.py b/test/diffsyn/raw_infer_data/create_input_output_npz.py new file mode 100644 index 00000000..072b850e --- /dev/null +++ b/test/diffsyn/raw_infer_data/create_input_output_npz.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python +"""Generate reference I/O arrays for DiffSyn regression testing. + +Equivalent PyTorch code (for cross-framework comparison): + import torch, numpy as np + torch.manual_seed(42); np.random.seed(42) + # Build identical DiffSyn (dim=32, channels=3, seq_length=8, + # cond_dim=16, dim_mults=(1,2), groups=4, timesteps=100, + # objective="pred_noise"), create dummy data, run forward pass. + # np.savez("reference_io.npz", pred=pred, loss=loss) + +This script uses the Paddle implementation to create a deterministic +reference snapshot. Re-run only when the model architecture intentionally +changes. +""" + +import importlib.util +import os +import sys + +import numpy as np +import paddle + +_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +_mod_path = os.path.join(_root, "ppmat", "models", "diffsyn", "diffsyn.py") +_spec = importlib.util.spec_from_file_location("diffsyn", _mod_path) +_diffsyn = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_diffsyn) +DiffSyn = _diffsyn.DiffSyn + + +def main(): + paddle.seed(42) + np.random.seed(42) + model = DiffSyn( + dim=32, channels=3, seq_length=8, cond_dim=16, + dim_mults=(1, 2), groups=4, timesteps=100, + objective="pred_noise", cond_drop_prob=0.0, + loss_type="l1_loss", + ) + model.eval() + + # Create deterministic input (re-seed for reproducible randn) + paddle.seed(42) + np.random.seed(42) + x = paddle.randn([2, 3, 8]) + cond = paddle.randn([2, 16]) + data = {"x": x, "cond": cond, "synthesis_conditions": x} + + result = model(data) + pred = result["pred_dict"]["synthesis_conditions"].numpy() + loss = result["loss_dict"]["loss"].numpy() + + out_path = os.path.join(os.path.dirname(__file__), "reference_io.npz") + np.savez(out_path, pred=pred, loss=loss) + print(f"Saved reference to {out_path}") + print(f" pred shape={pred.shape}") + print(f" loss={loss}") + + +if __name__ == "__main__": + main() diff --git a/test/diffsyn/raw_infer_data/reference_io.npz b/test/diffsyn/raw_infer_data/reference_io.npz new file mode 100644 index 00000000..53d55e95 Binary files /dev/null and b/test/diffsyn/raw_infer_data/reference_io.npz differ diff --git a/test/diffsyn/test_loss/__init__.py b/test/diffsyn/test_loss/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/diffsyn/test_loss/test_model_loss_with_raw.py b/test/diffsyn/test_loss/test_model_loss_with_raw.py new file mode 100644 index 00000000..5828e35a --- /dev/null +++ b/test/diffsyn/test_loss/test_model_loss_with_raw.py @@ -0,0 +1,313 @@ +"""Tests for DiffSyn model: shapes, loss, diffusion math, sampling.""" + +import importlib.util +import os +import sys +import unittest + +import numpy as np +import paddle + +# Load diffsyn.py directly to bypass ppmat's __init__ (which pulls pgl, etc.) +_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +_mod_path = os.path.join(_root, "ppmat", "models", "diffsyn", "diffsyn.py") +_spec = importlib.util.spec_from_file_location("diffsyn", _mod_path) +_diffsyn = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_diffsyn) + +DiffSyn = _diffsyn.DiffSyn +Unet1D = _diffsyn.Unet1D +SinusoidalPosEmb = _diffsyn.SinusoidalPosEmb +cosine_beta_schedule = _diffsyn.cosine_beta_schedule +extract = _diffsyn.extract + + +def _make_model(**overrides): + defaults = dict( + dim=32, channels=3, seq_length=8, cond_dim=16, + dim_mults=(1, 2), groups=4, timesteps=100, + objective="pred_noise", cond_drop_prob=0.0, + loss_type="l1_loss", + ) + defaults.update(overrides) + return DiffSyn(**defaults) + + +def _dummy_data(batch=2, channels=3, length=8, cond_dim=16): + x = paddle.randn([batch, channels, length]) + cond = paddle.randn([batch, cond_dim]) + return {"x": x, "cond": cond, "synthesis_conditions": x} + + +class TestDiffSynForwardAlignment(unittest.TestCase): + """Core forward-pass and shape tests.""" + + def setUp(self): + paddle.seed(42) + self.model = _make_model() + self.model.eval() + + def test_forward_returns_expected_keys(self): + data = _dummy_data() + out = self.model(data) + self.assertIn("loss_dict", out) + self.assertIn("pred_dict", out) + self.assertIn("loss", out["loss_dict"]) + self.assertIn("synthesis_conditions", out["pred_dict"]) + + def test_forward_prediction_shape(self): + data = _dummy_data() + out = self.model(data) + pred = out["pred_dict"]["synthesis_conditions"] + self.assertEqual(list(pred.shape), [2, 3, 8]) + + def test_loss_is_scalar(self): + data = _dummy_data() + out = self.model(data) + loss = out["loss_dict"]["loss"] + self.assertEqual(loss.ndim, 0) + self.assertFalse(paddle.isnan(loss).item()) + + def test_loss_only_mode(self): + data = _dummy_data() + out = self.model(data, return_loss=True, return_prediction=False) + self.assertIn("loss", out["loss_dict"]) + self.assertEqual(len(out["pred_dict"]), 0) + + def test_prediction_only_mode(self): + data = _dummy_data() + out = self.model(data, return_loss=False, return_prediction=True) + self.assertEqual(len(out["loss_dict"]), 0) + self.assertIn("synthesis_conditions", out["pred_dict"]) + + +class TestQSample(unittest.TestCase): + """Forward diffusion q(x_t | x_0) correctness.""" + + def setUp(self): + paddle.seed(0) + self.model = _make_model() + + def test_q_sample_t0_close_to_x_start(self): + """At t=0, x_t ≈ x_start (very little noise added).""" + x = paddle.randn([2, 3, 8]) + t = paddle.zeros([2], dtype="int64") + noise = paddle.randn(x.shape) + x_noisy = self.model.q_sample(x, t, noise) + # sqrt_alphas_cumprod[0] ≈ 1, so x_noisy ≈ x + tiny noise + diff = (x_noisy - x).abs().mean().item() + self.assertLess(diff, 0.15) + + def test_q_sample_reproducibility(self): + """Same seed → same noised sample.""" + x = paddle.ones([1, 3, 8]) + t = paddle.to_tensor([50], dtype="int64") + paddle.seed(99) + noise1 = paddle.randn(x.shape) + out1 = self.model.q_sample(x, t, noise1) + paddle.seed(99) + noise2 = paddle.randn(x.shape) + out2 = self.model.q_sample(x, t, noise2) + np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6) + + def test_q_sample_shape(self): + x = paddle.randn([4, 3, 8]) + t = paddle.randint(0, 100, [4]) + out = self.model.q_sample(x, t) + self.assertEqual(list(out.shape), [4, 3, 8]) + + +class TestDenoiseModel(unittest.TestCase): + """Direct Unet1D tests.""" + + def setUp(self): + paddle.seed(42) + self.unet = Unet1D(dim=32, channels=3, cond_dim=16, + dim_mults=(1, 2), groups=4, cond_drop_prob=0.0) + self.unet.eval() + + def test_unet_output_shape(self): + x = paddle.randn([2, 3, 8]) + t = paddle.randint(0, 100, [2]) + cond = paddle.randn([2, 16]) + out = self.unet(x, t, cond=cond) + self.assertEqual(list(out.shape), [2, 3, 8]) + + def test_unet_unconditional(self): + """Should work when cond=None (uses null embedding).""" + x = paddle.randn([2, 3, 8]) + t = paddle.randint(0, 100, [2]) + out = self.unet(x, t, cond=None) + self.assertEqual(list(out.shape), [2, 3, 8]) + + def test_cond_scale_forwarding(self): + x = paddle.randn([1, 3, 8]) + t = paddle.randint(0, 100, [1]) + cond = paddle.randn([1, 16]) + out_scale1 = self.unet.forward_with_cond_scale( + x, t, cond=cond, cond_scale=1.0) + out_scale3 = self.unet.forward_with_cond_scale( + x, t, cond=cond, cond_scale=3.0) + # scale=1 returns conditional; scale≠1 returns guided output + self.assertEqual(list(out_scale1.shape), list(out_scale3.shape)) + + +class TestSample(unittest.TestCase): + """Sampling (generative) path.""" + + def setUp(self): + paddle.seed(42) + self.model = _make_model(timesteps=10) + self.model.eval() + + def test_sample_shape(self): + samples = self.model.sample(batch_size=2, cond=paddle.randn([2, 16])) + self.assertEqual(list(samples.shape), [2, 3, 8]) + + def test_sample_unconditional(self): + samples = self.model.sample(batch_size=1) + self.assertEqual(list(samples.shape), [1, 3, 8]) + + def test_predict_method(self): + data = {"cond": paddle.randn([2, 16])} + result = self.model.predict(data) + self.assertIn("synthesis_conditions", result) + self.assertEqual(list(result["synthesis_conditions"].shape), [2, 3, 8]) + + +class TestNormalization(unittest.TestCase): + """Data normalize / unnormalize round-trip.""" + + def test_roundtrip(self): + model = _make_model(data_mean=2.0, data_std=3.0) + t = paddle.to_tensor([1.0, 5.0, 10.0]) + normed = model.normalize(t) + recovered = model.unnormalize(normed) + np.testing.assert_allclose(t.numpy(), recovered.numpy(), atol=1e-5) + + +class TestBetaSchedule(unittest.TestCase): + """Cosine beta schedule properties.""" + + def test_shape(self): + betas = cosine_beta_schedule(100) + self.assertEqual(betas.shape, (100,)) + + def test_range(self): + betas = cosine_beta_schedule(1000) + self.assertTrue(np.all(betas >= 0)) + self.assertTrue(np.all(betas <= 0.9999)) + + def test_monotonically_increasing_alphas_cumprod_decreases(self): + betas = cosine_beta_schedule(200) + alphas_cumprod = np.cumprod(1.0 - betas) + diffs = np.diff(alphas_cumprod) + self.assertTrue(np.all(diffs <= 0)) + + +class TestSinusoidalPosEmb(unittest.TestCase): + def test_output_shape(self): + emb = SinusoidalPosEmb(64) + t = paddle.to_tensor([0, 50, 99]) + out = emb(t) + self.assertEqual(list(out.shape), [3, 64]) + + +class TestExtract(unittest.TestCase): + def test_broadcast(self): + a = paddle.arange(10, dtype="float32") + t = paddle.to_tensor([3, 7], dtype="int64") + out = extract(a, t, [2, 4, 4]) + self.assertEqual(list(out.shape), [2, 1, 1]) + self.assertAlmostEqual(out[0, 0, 0].item(), 3.0) + self.assertAlmostEqual(out[1, 0, 0].item(), 7.0) + + +class TestObjectives(unittest.TestCase): + """Verify all three objective modes produce valid output.""" + + def test_pred_noise(self): + m = _make_model(objective="pred_noise") + out = m(_dummy_data()) + self.assertFalse(paddle.isnan(out["loss_dict"]["loss"]).item()) + + def test_pred_x0(self): + m = _make_model(objective="pred_x0") + out = m(_dummy_data()) + self.assertFalse(paddle.isnan(out["loss_dict"]["loss"]).item()) + + def test_pred_v(self): + m = _make_model(objective="pred_v") + out = m(_dummy_data()) + self.assertFalse(paddle.isnan(out["loss_dict"]["loss"]).item()) + + +class TestReferenceAlignment(unittest.TestCase): + """Hard-coded reference values for diffusion math reproducibility.""" + + def test_q_sample_known_values(self): + model = _make_model() + paddle.seed(123) + x = paddle.ones([1, 3, 8]) + t = paddle.to_tensor([0], dtype="int64") + noise = paddle.zeros([1, 3, 8]) + out = model.q_sample(x, t, noise) + # At t=0, sqrt_alphas_cumprod[0] ≈ 1.0, so out ≈ x + np.testing.assert_allclose(out.numpy(), x.numpy(), atol=0.01) + + def test_predict_start_from_noise_inverse(self): + """predict_start_from_noise should invert q_sample when noise is known.""" + model = _make_model() + x = paddle.randn([2, 3, 8]) + t = paddle.to_tensor([10, 50], dtype="int64") + noise = paddle.randn([2, 3, 8]) + x_noisy = model.q_sample(x, t, noise) + x_recovered = model.predict_start_from_noise(x_noisy, t, noise) + np.testing.assert_allclose( + x.numpy(), x_recovered.numpy(), atol=1e-4, rtol=1e-4 + ) + + +class TestTrainingConvergence(unittest.TestCase): + """Verify training convergence (RFC requirement).""" + + def test_training_convergence(self): + """Verify loss decreases over 2 training epochs (RFC requirement).""" + paddle.seed(42) + model = _make_model() + model.train() + optimizer = paddle.optimizer.Adam( + parameters=model.parameters(), learning_rate=1e-3 + ) + losses = [] + for epoch in range(2): + optimizer.clear_grad() + data = _dummy_data() + result = model(data) + loss = result["loss_dict"]["loss"] + loss.backward() + optimizer.step() + losses.append(loss.item()) + self.assertTrue( + np.isfinite(losses[-1]), f"Loss became non-finite: {losses}" + ) + self.assertLess( + losses[-1], + losses[0] * 1.5, + f"Loss increased significantly: {losses}", + ) + + +class TestMSELoss(unittest.TestCase): + """Verify MSE loss type works.""" + + def test_mse_loss(self): + model = _make_model(loss_type="mse_loss") + out = model(_dummy_data()) + loss = out["loss_dict"]["loss"] + self.assertFalse(paddle.isnan(loss).item()) + self.assertGreater(loss.item(), 0.0) + + +if __name__ == "__main__": + unittest.main()