-### 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 - -[](https://www.star-history.com/#PaddlePaddle/PaddleMaterilas&type=date&legend=top-left) - ---- +
| Parameter | +Description | +
|---|---|
| --model_name | +Name of the built-in model | +
| --weights_name | +Weights file name | +
| --cif_file_path | +Path to CIF files for prediction | +
| --save_path | +Path to save prediction results | +
- 
- 
- 
+ 
+ 
+ 
---- - ## 👩👩👧👦 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/newtonnet/README.md b/interatomic_potentials/configs/newtonnet/README.md new file mode 100644 index 00000000..26c2adcf --- /dev/null +++ b/interatomic_potentials/configs/newtonnet/README.md @@ -0,0 +1,48 @@ +# NewtonNet + +[NewtonNet: a Newtonian message passing network for deep learning of interatomic potentials and forces](https://doi.org/10.1039/D2DD00008C) + +## Abstract + +NewtonNet is a Newtonian message-passing neural network for learning interatomic potentials and forces. It operates on atomic graphs with radial Bessel basis functions and polynomial cutoff envelopes. The architecture uses a sequence of interaction layers that update both atom-level and force-level representations, enabling accurate prediction of energies and atomic forces while maintaining energy conservation. + +## Model + +NewtonNet constructs a radius graph from atomic positions and computes pairwise radial basis features. Interaction layers iteratively refine atom representations using Newtonian message passing that preserves physical symmetries. Per-atom energies are summed for the total molecular energy, and forces are obtained as negative energy gradients with respect to positions. + +## Training + +```bash +# single-gpu training +python interatomic_potentials/train.py -c interatomic_potentials/configs/newtonnet/newtonnet_qm9_energy.yaml + +# multi-gpu training +python -m paddle.distributed.launch --gpus="0,1,2,3" interatomic_potentials/train.py -c interatomic_potentials/configs/newtonnet/newtonnet_qm9_energy.yaml +``` + +## Validation + +```bash +python interatomic_potentials/train.py -c interatomic_potentials/configs/newtonnet/newtonnet_qm9_energy.yaml Global.do_eval=True Global.do_train=False Global.do_test=False Trainer.pretrained_model_path='your model path(*.pdparams)' +``` + +## Testing + +```bash +python interatomic_potentials/train.py -c interatomic_potentials/configs/newtonnet/newtonnet_qm9_energy.yaml Global.do_test=True Global.do_train=False Global.do_eval=False Trainer.pretrained_model_path='your model path(*.pdparams)' +``` + +## Citation + +``` +@article{haghighatlari2022newtonnet, + title={NewtonNet: a Newtonian message passing network for deep learning of interatomic potentials and forces}, + author={Haghighatlari, Mojtaba and Li, Jie and Guan, Xingyi and Zhang, Oufan and Das, Akshaya and Stein, Christopher J and Heiber, Farnaz and Liu, Tiantian and Head-Gordon, Martin and Bertels, Luke and others}, + journal={Digital Discovery}, + volume={1}, + number={3}, + pages={333--343}, + year={2022}, + publisher={Royal Society of Chemistry} +} +``` diff --git a/interatomic_potentials/configs/newtonnet/newtonnet_qm9_energy.yaml b/interatomic_potentials/configs/newtonnet/newtonnet_qm9_energy.yaml new file mode 100644 index 00000000..ab341b08 --- /dev/null +++ b/interatomic_potentials/configs/newtonnet/newtonnet_qm9_energy.yaml @@ -0,0 +1,89 @@ +Global: + do_train: True + do_eval: False + do_test: False + + label_names: ['energy'] + +Trainer: + max_epochs: 300 + seed: 42 + output_dir: ./output/newtonnet_qm9_energy + save_freq: 50 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + pretrained_model_path: null + pretrained_weight_name: null + resume_from_checkpoint: null + use_amp: False + amp_level: 'O1' + eval_with_no_grad: False + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_metric' + name_for_best_metric: "energy" + greater_is_better: False + compute_metric_during_train: True + metric_strategy_during_eval: 'epoch' + use_visualdl: False + use_wandb: False + use_tensorboard: False + +Model: + __class_name__: NewtonNet + __init_params__: + cutoff: 5.0 + n_features: 128 + n_basis: 20 + n_interactions: 3 + activation: swish + layer_norm: False + loss_type: mse_loss + force_loss_weight: 100.0 + property_names: ${Global.label_names} + +Optimizer: + __class_name__: Adam + __init_params__: + lr: + __class_name__: Cosine + __init_params__: + learning_rate: 5e-4 + eta_min: 1e-6 + by_epoch: False + +Metric: + energy: + __class_name__: paddle.nn.L1Loss + __init_params__: {} + +Dataset: + dataset: + __class_name__: QM9Dataset + __init_params__: + path: "./data/qm9" + property_names: ${Global.label_names} + num_workers: 4 + use_shared_memory: False + split_dataset_ratio: + train: 0.8 + val: 0.1 + test: 0.1 + train_sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val_sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test_sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 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..733154f1 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.newtonnet.newtonnet import NewtonNet from ppmat.utils import download from ppmat.utils import logger from ppmat.utils import save_load @@ -67,6 +68,7 @@ "DiffNMR", "InfGCN", "MatENO", + "NewtonNet", ] # Warning: The key of the dictionary must be consistent with the file name of the value diff --git a/ppmat/models/newtonnet/__init__.py b/ppmat/models/newtonnet/__init__.py new file mode 100644 index 00000000..675cf5cb --- /dev/null +++ b/ppmat/models/newtonnet/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/ppmat/models/newtonnet/newtonnet.py b/ppmat/models/newtonnet/newtonnet.py new file mode 100644 index 00000000..11fbab2f --- /dev/null +++ b/ppmat/models/newtonnet/newtonnet.py @@ -0,0 +1,743 @@ +# 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. + +""" +NewtonNet: Newtonian message passing network for molecular force fields. + +Adapted from the PyTorch implementation at: + https://github.com/THGLab/NewtonNet + +Reference: + Haghighatlari et al., "NewtonNet: a Newtonian message passing network + for deep learning of interatomic potentials and forces", 2022. +""" + +from typing import Optional + +import numpy as np +import paddle +import paddle.nn as nn + + +# --------------------------------------------------------------------------- +# Scatter utility (self-contained to avoid ppmat top-level import issues) +# Adapted from ppmat.utils.scatter +# --------------------------------------------------------------------------- + +def _broadcast(src: paddle.Tensor, other: paddle.Tensor, dim: int): + if dim < 0: + dim = other.ndim + dim + if src.ndim == 1: + for _ in range(0, dim): + src = src.unsqueeze(0) + for _ in range(src.ndim, other.ndim): + src = src.unsqueeze(-1) + src = src.expand(other.shape) + return src + + +def scatter( + src: paddle.Tensor, + index: paddle.Tensor, + dim: int = -1, + out: Optional[paddle.Tensor] = None, + dim_size: Optional[int] = None, + reduce: str = "sum", +) -> paddle.Tensor: + """Scatter-add (or mean) operation, compatible with torch_scatter.""" + index = _broadcast(index, src, dim) + if out is None: + size = list(src.shape) + if dim_size is not None: + size[dim] = dim_size + elif index.numel() == 0: + size[dim] = 0 + else: + size[dim] = int(index.max()) + 1 + out = paddle.zeros(size, dtype=src.dtype) + out = paddle.put_along_axis( + arr=out, indices=index, values=src, axis=dim, reduce="add" + ) + if reduce == "mean": + ones = paddle.ones(index.shape, dtype=src.dtype) + count = paddle.zeros(out.shape, dtype=src.dtype) + count = paddle.put_along_axis( + arr=count, indices=index, values=ones, axis=dim, reduce="add" + ) + count = paddle.clip(count, min=1) + out = out / count + return out + + +# --------------------------------------------------------------------------- +# Activation helpers +# --------------------------------------------------------------------------- + +def get_activation_by_string(key: str) -> nn.Layer: + """Return an activation layer from a string key.""" + activations = { + "swish": nn.Silu, + "silu": nn.Silu, + "relu": nn.ReLU, + "elu": nn.ELU, + "leaky_relu": nn.LeakyReLU, + "tanh": nn.Tanh, + "sigmoid": nn.Sigmoid, + "softplus": nn.Softplus, + "gelu": nn.GELU, + } + key_lower = key.lower() + if key_lower not in activations: + raise NotImplementedError(f"Unknown activation: {key}") + return activations[key_lower]() + + +# --------------------------------------------------------------------------- +# Representation layers +# --------------------------------------------------------------------------- + +class RadialBesselLayer(nn.Layer): + """Radial Bessel basis functions: ``sin(freq * dist) / dist``. + + Args: + n_basis (int): Number of radial basis functions. Default: 20. + """ + + def __init__(self, n_basis: int = 20): + super().__init__() + freqs = paddle.arange(1, n_basis + 1, dtype="float32") * float(np.pi) + self.register_buffer("frequencies", freqs) + + def forward(self, dist: paddle.Tensor) -> paddle.Tensor: + # dist: [E, 1] → output: [E, n_basis] + return paddle.sin(self.frequencies * dist) / dist + + +class PolynomialCutoff(nn.Layer): + """Smooth polynomial cutoff envelope. + + ``y = 1 - 0.5*(p+1)*(p+2)*x^p + p*(p+2)*x^(p+1) - 0.5*p*(p+1)*x^(p+2)`` + + Args: + p (int): Polynomial degree. Default: 9. + """ + + def __init__(self, p: int = 9): + super().__init__() + self.p = p + + def forward(self, dist: paddle.Tensor) -> paddle.Tensor: + p = self.p + x_p = dist ** p + x_p1 = x_p * dist + x_p2 = x_p1 * dist + return ( + 1.0 + - 0.5 * (p + 1) * (p + 2) * x_p + + p * (p + 2) * x_p1 + - 0.5 * p * (p + 1) * x_p2 + ) + + +class ScaledNorm(nn.Layer): + """Compute scaled distance and unit direction vector. + + Args: + r (float): Cutoff radius used for normalisation. + """ + + def __init__(self, r: float): + super().__init__() + self.r = r + + def forward(self, disp: paddle.Tensor): + """ + Args: + disp: Displacement vectors [E, 3]. + Returns: + dist: Scaled distances [E, 1]. + direction: Unit direction vectors [E, 3]. + """ + dist = paddle.norm(disp, axis=-1, keepdim=True) # [E, 1] + direction = disp / dist + dist = dist / self.r + return dist, direction + + +class RadiusGraph(nn.Layer): + """Build a radius graph from atomic positions with PBC support. + + Args: + r (float): Cutoff radius. + """ + + def __init__(self, r: float): + super().__init__() + self.r = r + + def forward(self, pos, cell=None, batch=None): + """ + Args: + pos: Atomic positions [N, 3]. + cell: Lattice vectors [B, 3, 3] or None. + batch: Batch indices [N] or None. + Returns: + edge_index: [2, E] source/target indices. + disp: Displacement vectors [E, 3]. + """ + n_node = pos.shape[0] + + # Build per-molecule full graph (excluding self-loops) + if batch is not None: + unique_batches = paddle.unique(batch) + src_list, dst_list = [], [] + for b in unique_batches: + mask = (batch == b) + nodes = paddle.nonzero(mask).flatten() + n = nodes.shape[0] + # Create all pairs + row = nodes.unsqueeze(1).expand([n, n]).flatten() + col = nodes.unsqueeze(0).expand([n, n]).flatten() + src_list.append(row) + dst_list.append(col) + src_all = paddle.concat(src_list) + dst_all = paddle.concat(dst_list) + edge_index = paddle.stack([src_all, dst_all], axis=0) + else: + idx = paddle.arange(n_node, dtype="int64") + row = idx.unsqueeze(1).expand([n_node, n_node]).flatten() + col = idx.unsqueeze(0).expand([n_node, n_node]).flatten() + edge_index = paddle.stack([row, col], axis=0) + + # Remove self-loops + non_self = edge_index[0] != edge_index[1] + edge_index = edge_index[:, non_self] + + # Displacement vectors + disp = pos[edge_index[0]] - pos[edge_index[1]] + + # Apply minimum-image convention for periodic boundary conditions + if cell is not None and not paddle.all(cell == 0.0): + if batch is not None: + cell_per_node = cell[batch] + else: + cell_per_node = cell.expand([n_node, 3, 3]) + cell_per_edge = cell_per_node[edge_index[0]] # [E, 3, 3] + # Solve for fractional coordinates: cell^T @ s = disp + cell_t = paddle.transpose(cell_per_edge, perm=[0, 2, 1]) # [E, 3, 3] + scaled = paddle.linalg.solve(cell_t, disp.unsqueeze(-1)).squeeze(-1) + # Detach before round: round() has zero gradient anyway, and + # solve backward segfaults on CPU in some Paddle versions. + disp = disp - paddle.bmm( + cell_per_edge, paddle.round(scaled.detach()).unsqueeze(-1) + ).squeeze(-1) + + # Filter by cutoff distance + dist_sq = (disp * disp).sum(axis=-1) + mask = dist_sq < self.r * self.r + edge_index = edge_index[:, mask] + disp = disp[mask] + + return edge_index, disp + + +class EdgeEmbedding(nn.Layer): + """Edge embedding combining RadiusGraph, ScaledNorm, PolynomialCutoff, and RadialBessel. + + Args: + cutoff (float): Cutoff radius. + n_basis (int): Number of radial basis functions. + """ + + def __init__(self, cutoff: float, n_basis: int = 20): + super().__init__() + self.radius_graph = RadiusGraph(r=cutoff) + self.norm = ScaledNorm(r=cutoff) + self.envelope = PolynomialCutoff(p=9) + self.embedding = RadialBesselLayer(n_basis=n_basis) + + def forward(self, pos, cell=None, batch=None): + edge_index, disp = self.radius_graph(pos, cell=cell, batch=batch) + dist_edge, dir_edge = self.norm(disp) + dist_edge = self.envelope(dist_edge) * self.embedding(dist_edge) + return dist_edge, dir_edge, edge_index + + +# --------------------------------------------------------------------------- +# Core network components +# --------------------------------------------------------------------------- + +class EmbeddingNet(nn.Layer): + """Initial atom and edge embedding layer. + + Args: + cutoff (float): Cutoff radius for edge embedding. + n_features (int): Feature dimension. + n_basis (int): Number of radial basis functions. + """ + + def __init__(self, cutoff: float, n_features: int, n_basis: int): + super().__init__() + self.n_features = n_features + self.node_embedding = nn.Embedding(118 + 1, n_features, padding_idx=0) + self.edge_embedding = EdgeEmbedding(cutoff=cutoff, n_basis=n_basis) + self.requires_dr = False + + def forward(self, z, pos, cell, batch): + # Node embedding + atom_node = self.node_embedding(z) # [N, F] + force_node = paddle.zeros( + [z.shape[0], 3, self.n_features], dtype=pos.dtype + ) # [N, 3, F] + + # Strain displacement for virial/stress (identity by default) + displacement = paddle.zeros_like(cell) + displacement[:, 0, 0] = 1.0 + displacement[:, 1, 1] = 1.0 + displacement[:, 2, 2] = 1.0 + + if self.requires_dr: + pos.stop_gradient = False + displacement.stop_gradient = False + + symmetric_displacement = ( + displacement + paddle.transpose(displacement, perm=[0, 2, 1]) + ) / 2.0 + # Apply strain to positions: pos_displaced = pos @ symmetric_displacement + pos_displaced = paddle.bmm( + pos.unsqueeze(1), symmetric_displacement[batch] + ).squeeze(1) # [N, 3] + cell_displaced = paddle.bmm( + cell, symmetric_displacement + ).squeeze(1) # [B, 3, 3] (squeeze does nothing if already [B,3,3]) + + # Edge embedding + dist_edge, dir_edge, edge_index = self.edge_embedding( + pos_displaced, cell_displaced, batch + ) + + return atom_node, force_node, dir_edge, dist_edge, edge_index, displacement + + +class InteractionNet(nn.Layer): + """Newtonian message passing interaction layer. + + Performs both invariant (scalar) and equivariant (3D vector) message passing. + + Args: + n_features (int): Feature dimension. + n_basis (int): Number of radial basis functions. + activation (nn.Layer): Activation function instance. + layer_norm (bool): Whether to apply LayerNorm to atom features. + """ + + def __init__( + self, + n_features: int, + n_basis: int, + activation: nn.Layer, + layer_norm: bool, + ): + super().__init__() + self.n_features = n_features + + # Invariant message passing + self.message_nodepart = nn.Sequential( + nn.Linear(n_features, n_features), + activation, + nn.Linear(n_features, n_features), + ) + self.message_edgepart = nn.Linear(n_basis, n_features, bias_attr=False) + + # Equivariant message networks + self.equiv_message1 = nn.Sequential( + nn.Linear(n_features, n_features, bias_attr=False), + activation, + nn.Linear(n_features, n_features, bias_attr=False), + ) + self.equiv_message2 = nn.Sequential( + nn.Linear(n_features, n_features, bias_attr=False), + activation, + nn.Linear(n_features, n_features, bias_attr=False), + ) + + # Force → energy update + self.equiv_update = nn.Linear(n_features, n_features, bias_attr=False) + + # Optional layer norm + self.layer_norm = nn.LayerNorm(n_features) if layer_norm else None + + def forward(self, atom_node, force_node, dir_edge, dist_edge, edge_index): + """ + Args: + atom_node: [N, F] + force_node: [N, 3, F] + dir_edge: [E, 3] + dist_edge: [E, n_basis] + edge_index: [2, E] (row=src, col=dst) + Returns: + Updated atom_node [N, F] and force_node [N, 3, F]. + """ + src, dst = edge_index[0], edge_index[1] + + # --- Invariant message --- + message_nodepart = self.message_nodepart(atom_node) # [N, F] + message_edgepart = self.message_edgepart(dist_edge) # [E, F] + message = ( + message_edgepart + * message_nodepart[src] + * message_nodepart[dst] + ) # [E, F] + + inv_update = scatter( + message, src, dim=0, dim_size=atom_node.shape[0] + ) # [N, F] + atom_node = atom_node + inv_update + + # --- Equivariant message 1: direction-weighted --- + equiv_msg1_inv = self.equiv_message1(message).unsqueeze(1) # [E, 1, F] + equiv_msg1_eq = dir_edge.unsqueeze(2) # [E, 3, 1] + equiv_msg1 = equiv_msg1_inv * equiv_msg1_eq # [E, 3, F] + + # --- Equivariant message 2: neighbor force --- + equiv_msg2_inv = self.equiv_message2(message).unsqueeze(1) # [E, 1, F] + equiv_msg2_eq = force_node[dst] # [E, 3, F] + equiv_msg2 = equiv_msg2_inv * equiv_msg2_eq # [E, 3, F] + + force_update = scatter( + equiv_msg1 + equiv_msg2, src, dim=0, dim_size=force_node.shape[0] + ) # [N, 3, F] + force_node = force_node + force_update + + # --- Energy update from force (dot product over xyz) --- + inv_update2 = paddle.sum( + force_node * self.equiv_update(force_node), axis=1 + ) # [N, F] + atom_node = atom_node + inv_update2 + + # --- Layer norm --- + if self.layer_norm is not None: + atom_node = self.layer_norm(atom_node) + + return atom_node, force_node + + +# --------------------------------------------------------------------------- +# Output layers +# --------------------------------------------------------------------------- + +class EnergyOutput(nn.Layer): + """Per-atom energy prediction via MLP. + + Args: + n_features (int): Feature dimension. + activation (nn.Layer): Activation function instance. + """ + + def __init__(self, n_features: int, activation: nn.Layer): + super().__init__() + self.layers = nn.Sequential( + nn.Linear(n_features, n_features), + activation, + nn.Linear(n_features, n_features), + activation, + nn.Linear(n_features, 1), + ) + + def forward(self, atom_node: paddle.Tensor) -> paddle.Tensor: + """Return per-atom energy [N, 1].""" + return self.layers(atom_node) + + +class DirectForceOutput(nn.Layer): + """Direct force prediction from atom features and equivariant force node. + + Args: + n_features (int): Feature dimension. + activation (nn.Layer): Activation function instance. + """ + + def __init__(self, n_features: int, activation: nn.Layer): + super().__init__() + self.layers = nn.Sequential( + nn.Linear(n_features, n_features), + activation, + nn.Linear(n_features, n_features), + activation, + nn.Linear(n_features, n_features), + ) + + def forward( + self, atom_node: paddle.Tensor, force_node: paddle.Tensor + ) -> paddle.Tensor: + """Return forces [N, 3].""" + weight = self.layers(atom_node).unsqueeze(1) # [N, 1, F] + force = weight * force_node # [N, 3, F] + return force.sum(axis=-1) # [N, 3] + + +class ScaleShift(nn.Layer): + """Per-element scale and shift using atomic-number embeddings. + + Args: + scale (bool): Whether to apply scaling. + shift (bool): Whether to apply shifting. + """ + + def __init__(self, scale: bool = True, shift: bool = True): + super().__init__() + if scale: + self.scale = nn.Embedding(119, 1, padding_idx=0) + # Initialize scale to 1 + with paddle.no_grad(): + self.scale.weight.set_value( + paddle.ones_like(self.scale.weight) + ) + else: + self.scale = None + if shift: + self.shift = nn.Embedding(119, 1, padding_idx=0) + # Initialize shift to 0 (default) + else: + self.shift = None + + def forward(self, output, z): + if self.scale is not None: + output = output * self.scale(z) + if self.shift is not None: + output = output + self.shift(z) + return output + + +# --------------------------------------------------------------------------- +# Main model +# --------------------------------------------------------------------------- + +class NewtonNet(nn.Layer): + """Newtonian message passing network for molecular force fields. + + Follows PaddleMaterials conventions with ``forward()`` returning + ``{"loss_dict": ..., "pred_dict": ...}``. + + Args: + cutoff (float): Cutoff radius (Å). Default: 5.0. + n_features (int): Hidden feature dimension. Default: 128. + n_basis (int): Number of radial basis functions. Default: 20. + n_interactions (int): Number of message passing layers. Default: 3. + activation (str): Activation function name. Default: ``"swish"``. + layer_norm (bool): Apply LayerNorm after each interaction. Default: False. + property_names (str): Target property name. Default: ``"energy"``. + data_mean (float): Mean for target normalisation. Default: 0.0. + data_std (float): Std for target normalisation. Default: 1.0. + loss_type (str): ``"mse_loss"`` or ``"l1_loss"``. Default: ``"mse_loss"``. + force_loss_weight (float): Weight for force loss relative to energy loss. + Default: 100.0. + """ + + def __init__( + self, + cutoff: float = 5.0, + n_features: int = 128, + n_basis: int = 20, + n_interactions: int = 3, + activation: str = "swish", + layer_norm: bool = False, + property_names: str = "energy", + data_mean: float = 0.0, + data_std: float = 1.0, + loss_type: str = "mse_loss", + force_loss_weight: float = 100.0, + ): + super().__init__() + + if isinstance(property_names, list): + self.property_names = property_names[0] + else: + self.property_names = property_names + + self.register_buffer("data_mean", paddle.to_tensor(data_mean, dtype="float32")) + self.register_buffer("data_std", paddle.to_tensor(data_std, dtype="float32")) + + act = get_activation_by_string(activation) + + # Embedding + self.embedding_layer = EmbeddingNet( + cutoff=cutoff, + n_features=n_features, + n_basis=n_basis, + ) + + # Interaction layers + self.interaction_layers = nn.LayerList( + [ + InteractionNet( + n_features=n_features, + n_basis=n_basis, + activation=get_activation_by_string(activation), + layer_norm=layer_norm, + ) + for _ in range(n_interactions) + ] + ) + + # Energy output + self.energy_output = EnergyOutput( + n_features, get_activation_by_string(activation) + ) + + # Per-element energy scale/shift + self.energy_scaler = ScaleShift(scale=True, shift=True) + + # Loss + self.force_loss_weight = force_loss_weight + + # Loss + if loss_type == "mse_loss": + self.loss_fn = nn.functional.mse_loss + elif loss_type == "l1_loss": + self.loss_fn = nn.functional.l1_loss + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + # ------------------------------------------------------------------ + # Normalisation helpers + # ------------------------------------------------------------------ + def normalize(self, tensor: paddle.Tensor) -> paddle.Tensor: + return (tensor - self.data_mean) / self.data_std + + def unnormalize(self, tensor: paddle.Tensor) -> paddle.Tensor: + return tensor * self.data_std + self.data_mean + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + def _forward(self, data: dict, compute_forces: bool = False): + """Core forward pass returning per-molecule energy and optional forces. + + Args: + data: Dict with keys ``z``, ``pos``, ``batch``, ``cell`` (optional). + compute_forces: If True, compute forces as negative energy gradient. + Returns: + Tuple of (energy [B], forces [N, 3] or None). + """ + z = data["z"] + pos = data["pos"] + batch = data["batch"] + cell = data.get("cell", None) + + if compute_forces: + pos.stop_gradient = False + + if cell is None: + n_batch = int(batch.max().item()) + 1 + cell = paddle.zeros([n_batch, 3, 3], dtype=pos.dtype) + + # Embedding + atom_node, force_node, dir_edge, dist_edge, edge_index, displacement = ( + self.embedding_layer(z, pos, cell, batch) + ) + + # Message passing + for interaction in self.interaction_layers: + atom_node, force_node = interaction( + atom_node, force_node, dir_edge, dist_edge, edge_index + ) + + # Per-atom energy + energy_per_atom = self.energy_output(atom_node) # [N, 1] + energy_per_atom = self.energy_scaler(energy_per_atom, z) # [N, 1] + + # Compute forces as F = -dE/dpos (in physical units) + forces = None + if compute_forces: + grad_result = paddle.grad( + outputs=energy_per_atom.sum(), + inputs=pos, + create_graph=self.training, + retain_graph=True, + )[0] + forces = -grad_result * self.data_std + + # Aggregate per molecule + energy = scatter( + energy_per_atom, batch, dim=0, dim_size=int(batch.max().item()) + 1 + ).reshape([-1]) # [B] + + return energy, forces + + def forward( + self, + data: dict, + return_loss: bool = True, + return_prediction: bool = True, + ) -> dict: + """PaddleMaterials standard forward. + + Args: + data: Input data dict. + return_loss: Whether to compute loss. + return_prediction: Whether to return predictions. + Returns: + Dict with ``loss_dict`` and ``pred_dict``. + """ + assert ( + return_loss or return_prediction + ), "At least one of return_loss or return_prediction must be True." + + compute_forces = "forces" in data or return_prediction + energy, forces = self._forward(data, compute_forces=compute_forces) + + loss_dict = {} + if return_loss: + label = data[self.property_names] + label = self.normalize(label) + energy_loss = self.loss_fn(input=energy, label=label) + + total_loss = energy_loss + if "forces" in data and forces is not None: + force_loss = self.loss_fn(input=forces, label=data["forces"]) + loss_dict["force_loss"] = force_loss + total_loss = energy_loss + self.force_loss_weight * force_loss + loss_dict["loss"] = total_loss + + prediction = {} + if return_prediction: + prediction[self.property_names] = self.unnormalize(energy) + if forces is not None: + prediction["forces"] = forces + + return {"loss_dict": loss_dict, "pred_dict": prediction} + + def predict(self, data: dict, compute_forces: bool = True) -> dict: + """Predict energy and optionally forces. + + Args: + data: Input data dict. + compute_forces: Whether to compute forces. Default: True. + Returns: + Dict mapping property name to predicted value, plus ``"forces"`` + when *compute_forces* is True. + """ + if compute_forces: + energy, forces = self._forward(data, compute_forces=True) + energy = self.unnormalize(energy) + result = {self.property_names: energy.detach()} + if forces is not None: + result["forces"] = forces.detach() + return result + else: + with paddle.no_grad(): + energy, _ = self._forward(data, compute_forces=False) + energy = self.unnormalize(energy) + return {self.property_names: energy} diff --git a/test/newtonnet/__init__.py b/test/newtonnet/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/newtonnet/conftest.py b/test/newtonnet/conftest.py new file mode 100644 index 00000000..a58a718f --- /dev/null +++ b/test/newtonnet/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 NewtonNet 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/newtonnet/raw_infer_data/__init__.py b/test/newtonnet/raw_infer_data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/newtonnet/raw_infer_data/create_input_output_npz.py b/test/newtonnet/raw_infer_data/create_input_output_npz.py new file mode 100644 index 00000000..c18bd567 --- /dev/null +++ b/test/newtonnet/raw_infer_data/create_input_output_npz.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python +"""Generate reference I/O arrays for NewtonNet regression testing. + +Equivalent PyTorch code (for cross-framework comparison): + import torch, numpy as np + torch.manual_seed(42); np.random.seed(42) + # Build identical NewtonNet with PyTorch (cutoff=5.0, n_features=32, + # n_basis=8, n_interactions=2, activation="swish"), create water + # molecule input (H2O), run forward pass, save energy output. + # np.savez("reference_io.npz", energy=energy, 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 numpy as np +import paddle + +_module_path = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", "..", "..", + "ppmat", "models", "newtonnet", "newtonnet.py", + ) +) +_spec = importlib.util.spec_from_file_location("newtonnet_module", _module_path) +_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mod) +NewtonNet = _mod.NewtonNet + + +def main(): + paddle.seed(42) + np.random.seed(42) + model = NewtonNet( + cutoff=5.0, n_features=32, n_basis=8, + n_interactions=2, activation="swish", + layer_norm=False, property_names="energy", + data_mean=0.0, data_std=1.0, loss_type="mse_loss", + ) + model.eval() + + data = { + "z": paddle.to_tensor([8, 1, 1], dtype="int64"), + "pos": paddle.to_tensor( + [[0.0, 0.0, 0.1173], + [0.0, 0.7572, -0.4692], + [0.0, -0.7572, -0.4692]], + dtype="float32", + ), + "batch": paddle.zeros([3], dtype="int64"), + "cell": paddle.eye(3, dtype="float32").unsqueeze(0) * 10.0, + "energy": paddle.to_tensor([-76.4], dtype="float32"), + } + + result = model(data, return_loss=True, return_prediction=True) + energy = result["pred_dict"]["energy"].numpy() + loss = result["loss_dict"]["loss"].numpy() + + out_path = os.path.join(os.path.dirname(__file__), "reference_io.npz") + np.savez(out_path, energy=energy, loss=loss) + print(f"Saved reference to {out_path}") + print(f" energy shape={energy.shape}, values={energy}") + print(f" loss={loss}") + + +if __name__ == "__main__": + main() diff --git a/test/newtonnet/raw_infer_data/reference_io.npz b/test/newtonnet/raw_infer_data/reference_io.npz new file mode 100644 index 00000000..959b29a5 Binary files /dev/null and b/test/newtonnet/raw_infer_data/reference_io.npz differ diff --git a/test/newtonnet/test_loss/__init__.py b/test/newtonnet/test_loss/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/newtonnet/test_loss/test_model_loss_with_raw.py b/test/newtonnet/test_loss/test_model_loss_with_raw.py new file mode 100644 index 00000000..5fb515fb --- /dev/null +++ b/test/newtonnet/test_loss/test_model_loss_with_raw.py @@ -0,0 +1,262 @@ +# 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. + +"""Tests for NewtonNet forward pass, loss computation, and prediction structure.""" + +import importlib.util +import os +import unittest + +import numpy as np +import paddle + +# Load the module directly from its file path to avoid ppmat/__init__.py +# which requires pgl (PaddleGraphLearning) not available in this env. +_module_path = os.path.join( + os.path.dirname(__file__), + "..", "..", "..", + "ppmat", "models", "newtonnet", "newtonnet.py", +) +_module_path = os.path.abspath(_module_path) +_spec = importlib.util.spec_from_file_location("newtonnet_module", _module_path) +_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mod) +NewtonNet = _mod.NewtonNet + + +class TestNewtonNetForwardAlignment(unittest.TestCase): + """Test NewtonNet forward pass alignment and basic functionality.""" + + def setUp(self): + paddle.seed(42) + self.model = NewtonNet( + cutoff=5.0, + n_features=32, + n_basis=8, + n_interactions=2, + activation="swish", + layer_norm=False, + property_names="energy", + data_mean=0.0, + data_std=1.0, + loss_type="mse_loss", + ) + self.model.eval() + + def _create_water_molecule_data(self): + """Create a single water molecule (H2O) test input.""" + z = paddle.to_tensor([8, 1, 1], dtype="int64") # O, H, H + pos = paddle.to_tensor( + [ + [0.0000, 0.0000, 0.1173], + [0.0000, 0.7572, -0.4692], + [0.0000, -0.7572, -0.4692], + ], + dtype="float32", + ) + batch = paddle.zeros([3], dtype="int64") + cell = paddle.eye(3, dtype="float32").unsqueeze(0) * 10.0 + energy = paddle.to_tensor([-76.4], dtype="float32") + return { + "z": z, + "pos": pos, + "batch": batch, + "cell": cell, + "energy": energy, + } + + def _create_two_molecule_batch(self): + """Create a batch with two molecules: H2O and H2.""" + z = paddle.to_tensor([8, 1, 1, 1, 1], dtype="int64") + pos = paddle.to_tensor( + [ + [0.0, 0.0, 0.1173], + [0.0, 0.7572, -0.4692], + [0.0, -0.7572, -0.4692], + [5.0, 0.0, 0.0], + [5.0, 0.74, 0.0], + ], + dtype="float32", + ) + batch = paddle.to_tensor([0, 0, 0, 1, 1], dtype="int64") + cell = paddle.eye(3, dtype="float32").unsqueeze(0).expand([2, 3, 3]) * 10.0 + energy = paddle.to_tensor([-76.4, -1.17], dtype="float32") + return { + "z": z, + "pos": pos, + "batch": batch, + "cell": cell, + "energy": energy, + } + + def test_forward_shape(self): + """Verify output shapes from forward pass.""" + data = self._create_water_molecule_data() + result = self.model(data, return_loss=True, return_prediction=True) + + self.assertIn("loss_dict", result) + self.assertIn("pred_dict", result) + self.assertIn("loss", result["loss_dict"]) + self.assertIn("energy", result["pred_dict"]) + + loss = result["loss_dict"]["loss"] + energy = result["pred_dict"]["energy"] + self.assertEqual(loss.shape, []) # scalar + self.assertEqual(energy.shape, [1]) # one molecule + + def test_forward_shape_batch(self): + """Verify shapes for a two-molecule batch.""" + data = self._create_two_molecule_batch() + result = self.model(data, return_loss=True, return_prediction=True) + + energy = result["pred_dict"]["energy"] + self.assertEqual(energy.shape, [2]) + + def test_forward_determinism(self): + """Two forward passes with the same input must produce identical results.""" + data = self._create_water_molecule_data() + r1 = self.model(data, return_loss=False, return_prediction=True) + r2 = self.model(data, return_loss=False, return_prediction=True) + np.testing.assert_allclose( + r1["pred_dict"]["energy"].numpy(), + r2["pred_dict"]["energy"].numpy(), + rtol=1e-6, + ) + + def test_loss_computation(self): + """Loss should be a finite positive scalar.""" + data = self._create_water_molecule_data() + result = self.model(data, return_loss=True, return_prediction=False) + + loss = result["loss_dict"]["loss"] + self.assertEqual(loss.shape, []) + self.assertTrue(paddle.isfinite(loss).item()) + self.assertGreaterEqual(loss.item(), 0.0) + + def test_energy_prediction_structure(self): + """Predicted energy must be finite and return the correct key.""" + data = self._create_water_molecule_data() + result = self.model(data, return_loss=False, return_prediction=True) + + self.assertIn("energy", result["pred_dict"]) + energy = result["pred_dict"]["energy"] + self.assertTrue(paddle.isfinite(energy).all().item()) + + def test_predict_method(self): + """The convenience ``predict()`` method should return a dict with energy.""" + data = self._create_water_molecule_data() + pred = self.model.predict(data) + self.assertIn("energy", pred) + self.assertEqual(pred["energy"].shape, [1]) + self.assertTrue(paddle.isfinite(pred["energy"]).all().item()) + + def test_force_computation(self): + """Forces are computed as negative energy gradient.""" + data = self._create_water_molecule_data() + data["pos"].stop_gradient = False + result = self.model(data, return_loss=False, return_prediction=True) + forces = result["pred_dict"]["forces"] + self.assertEqual(list(forces.shape), [data["pos"].shape[0], 3]) + self.assertTrue(paddle.isfinite(forces).all().item()) + + def test_energy_force_consistency(self): + """Forces should approximate the negative finite-difference gradient.""" + data = self._create_water_molecule_data() + pos = data["pos"].clone() + pos.stop_gradient = False + data["pos"] = pos + result = self.model(data, return_loss=False, return_prediction=True) + pred_forces = result["pred_dict"]["forces"] + + # Verify via finite differences + eps = 1e-3 + n_atoms = pos.shape[0] + fd_forces = paddle.zeros([n_atoms, 3], dtype="float32") + for i in range(n_atoms): + for j in range(3): + data_p = self._create_water_molecule_data() + data_p["pos"] = data_p["pos"].clone() + data_p["pos"][i, j] += eps + e_p = self.model( + data_p, return_loss=False, return_prediction=True + )["pred_dict"]["energy"] + + data_m = self._create_water_molecule_data() + data_m["pos"] = data_m["pos"].clone() + data_m["pos"][i, j] -= eps + e_m = self.model( + data_m, return_loss=False, return_prediction=True + )["pred_dict"]["energy"] + + fd_forces[i, j] = -(e_p.sum() - e_m.sum()) / (2 * eps) + + np.testing.assert_allclose( + pred_forces.detach().numpy(), fd_forces.numpy(), atol=5e-2 + ) + + def test_force_loss(self): + """Forward pass with force targets computes force loss.""" + data = self._create_water_molecule_data() + data["forces"] = paddle.randn(data["pos"].shape) + data["pos"].stop_gradient = False + result = self.model(data, return_loss=True, return_prediction=True) + self.assertIn("force_loss", result["loss_dict"]) + self.assertTrue(result["loss_dict"]["force_loss"].item() > 0) + + def test_normalize_unnormalize_roundtrip(self): + """normalize → unnormalize should recover the original value.""" + model = NewtonNet(data_mean=-76.0, data_std=10.0) + original = paddle.to_tensor([-76.4], dtype="float32") + normalized = model.normalize(original) + recovered = model.unnormalize(normalized) + np.testing.assert_allclose( + recovered.numpy(), original.numpy(), rtol=1e-5 + ) + + def test_no_cell_input(self): + """Model should work when cell is not provided.""" + data = self._create_water_molecule_data() + del data["cell"] + result = self.model(data, return_loss=False, return_prediction=True) + energy = result["pred_dict"]["energy"] + self.assertEqual(energy.shape, [1]) + self.assertTrue(paddle.isfinite(energy).all().item()) + + def test_training_convergence(self): + """Verify loss decreases over 2 training epochs (RFC requirement).""" + self.model.train() + optimizer = paddle.optimizer.Adam( + parameters=self.model.parameters(), learning_rate=1e-3 + ) + losses = [] + for epoch in range(2): + optimizer.clear_grad() + data = self._create_water_molecule_data() + result = self.model(data, return_loss=True, return_prediction=False) + 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}", + ) + + +if __name__ == "__main__": + unittest.main()