diff --git a/.gitignore b/.gitignore index e4dd3a77..188a8044 100644 --- a/.gitignore +++ b/.gitignore @@ -180,3 +180,4 @@ electronic_structure/configs/qm9_data_split.json electronic_structure/configs/qm9.json electronic_structure/configs/crystal_data_split.json electronic_structure/configs/crystal.json +pretrained_weights/trained_schnet_models/ diff --git a/README.md b/README.md index ddcc5c78..d0bf9e87 100755 --- a/README.md +++ b/README.md @@ -6,176 +6,127 @@ ## 🚀 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: +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. + + -[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 + diff --git a/checkpoints/schnet_md17_ethanol/schnet_md17_ethanol.pdparams b/checkpoints/schnet_md17_ethanol/schnet_md17_ethanol.pdparams new file mode 100644 index 00000000..e4c76ad0 Binary files /dev/null and b/checkpoints/schnet_md17_ethanol/schnet_md17_ethanol.pdparams differ diff --git a/checkpoints/schnet_qm9_U0/schnet_qm9_U0.pdparams b/checkpoints/schnet_qm9_U0/schnet_qm9_U0.pdparams new file mode 100644 index 00000000..de7b8d0b Binary files /dev/null and b/checkpoints/schnet_qm9_U0/schnet_qm9_U0.pdparams differ 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/README.md b/interatomic_potentials/README.md index ae0cd3b2..398f7659 100644 --- a/interatomic_potentials/README.md +++ b/interatomic_potentials/README.md @@ -6,30 +6,32 @@ Machine-learning interatomic potentials (MLIP) bridge the gap between quantum-le ## 2.Models Matrix -| **Supported Functions** | **[CHGNet](./configs/chgnet/README.md)** | **[MatterSim](./configs/mattersim//README.md)** | -| ----------------------------------- | ---------------------------------------- | ----------------------------------------------- | -| **Forward Prediction** | | | -|  Energy | ✅ | ✅ | -|  Force | ✅ | ✅ | -|  Stress | ✅ | ✅ | -|  Magmom | ✅ | - | -| **ML Capabilities · Training** | | | -|  Single-GPU | ✅ | ✅ | -|  Distributed Train | ✅ | ✅ | -|  Mixed Precision | - | - | -|  Fine-tuning | ✅ | ✅ | -|  Uncertainty / Active-Learning | - | - | -|  Dynamic→Static | - | - | -|  Compiler CINN | - | - | -| **ML Capabilities · Predict** | | | -|  Distillation / Pruning | - | - | -|  Standard inference | ✅ | ✅ | -|  Distributed inference | - | - | -|  Compiler CINN | - | - | -| **Molecular Dynamic Interface** | | | -|  ASE | ✅ | ✅ | -| **Dataset** | | | -|  MPtrj | ✅ | 🚧 | -| **ML2DDB🌟** | ✅ | - | +| **Supported Functions** | **[CHGNet](./configs/chgnet/README.md)** | **[MatterSim](./configs/mattersim//README.md)** | **[SchNet](./configs/schnet/README.md)** | +| ----------------------------------- | ---------------------------------------- | ----------------------------------------------- | ---------------------------------------- | +| **Forward Prediction** | | | | +|  Energy | ✅ | ✅ | ✅ | +|  Force | ✅ | ✅ | - | +|  Stress | ✅ | ✅ | - | +|  Magmom | ✅ | - | - | +| **ML Capabilities · Training** | | | | +|  Single-GPU | ✅ | ✅ | ✅ | +|  Distributed Train | ✅ | ✅ | ✅ | +|  Mixed Precision | - | - | - | +|  Fine-tuning | ✅ | ✅ | ✅ | +|  Uncertainty / Active-Learning | - | - | - | +|  Dynamic→Static | - | - | - | +|  Compiler CINN | - | - | - | +| **ML Capabilities · Predict** | | | | +|  Distillation / Pruning | - | - | - | +|  Standard inference | ✅ | ✅ | ✅ | +|  Distributed inference | - | - | - | +|  Compiler CINN | - | - | - | +| **Molecular Dynamic Interface** | | | | +|  ASE | ✅ | ✅ | - | +| **Dataset** | | | | +|  MPtrj | ✅ | 🚧 | - | +|  QM9 | - | - | ✅ | +|  MD17 | - | - | ✅ | +| **ML2DDB🌟** | ✅ | - | - | **Notice**:🌟 represent originate research work published from paddlematerials toolkit diff --git a/interatomic_potentials/configs/schnet/README.md b/interatomic_potentials/configs/schnet/README.md new file mode 100644 index 00000000..24a9dd9a --- /dev/null +++ b/interatomic_potentials/configs/schnet/README.md @@ -0,0 +1,159 @@ +# SchNet + +[SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://proceedings.neurips.cc/paper_files/paper/2017/hash/303ed4c69846ab36c2904d3ba8573050-Abstract.html) + +## Abstract + +Deep learning has the potential to revolutionize quantum chemistry as it is ideally suited to learn representations for structured data and speed up the exploration of chemical compound space. While convolutional neural networks have proven to be the first choice for images, audio and video data, the atoms in molecules are not restricted to a grid. Instead, their precise locations contain essential physical information, that would get lost if discretized. We propose to use continuous-filter convolutional layers to be able to model local correlations without requiring the data to lie on a grid. SchNet achieves state-of-the-art accuracy for predicting total energies and forces of molecules with DFT accuracy, while being several orders of magnitude faster than DFT calculations. + +![SchNet Overview](../../docs/schnet.png) + +## Datasets + +SchNet is trained and evaluated on two benchmark datasets for molecular property prediction: + +- **QM9**: + + The QM9 dataset is a widely-used benchmark for predicting molecular properties of small organic molecules. It contains 133,885 small organic molecules with up to 9 heavy atoms (C, H, O, N, F). Properties are computed at the B3LYP/6-31G(2df,p) level of DFT. + + We use the standard data split following SchNetPack convention: + | Dataset | Train | Val | Test | + | :------------------------------------------: | :----: | :----: | :---: | + | QM9 | 100000 | 10000 | 23885 | + + The target property for the QM9 configuration is the internal energy at 0K (U0). + +- **MD17** (Ethanol): + + The MD17 dataset contains molecular dynamics trajectories computed at the PBE+vdW-TS level of DFT. The ethanol molecule trajectory is used here for energy prediction. + + | Dataset | Train | Val | + | :----------------: | :----: | :----: | + | md17_ethanol | 50000 | 10000 | + + +## Models + +SchNet uses a message-passing architecture based on continuous-filter convolutions (CFConv). Given atomic coordinates, SchNet operates in three stages: + +### 1. Atom Embedding +Each atom is mapped to a learnable feature vector based on its atomic number: +$$ +\mathbf{x}_i^{(0)} = \mathbf{a}_{Z_i} +$$ + +### 2. Interaction Blocks +Distance-dependent filters are generated by expanding interatomic distances using a Gaussian radial basis: +$$ +e_k(r_{ij}) = \exp\left(-\gamma\left(r_{ij} - \mu_k\right)^2\right) +$$ + +These are passed through a filter-generating network containing a shifted softplus activation: +$$ +\text{ssp}(x) = \ln(1 + e^x) - \ln 2 +$$ + +The resulting continuous filters are used to compute atom-wise convolutions: +$$ +\mathbf{x}_i^{(t+1)} = \mathbf{x}_i^{(t)} + \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j^{(t)} \circ W^{(t)}(\mathbf{e}_{ij}) +$$ + +where $\circ$ denotes element-wise multiplication and $W^{(t)}$ is the filter-generating network at interaction $t$. + +### 3. Output Network +Atom-wise features are mapped to the target property through a two-layer MLP with shifted softplus activation, followed by a sum aggregation over atoms: +$$ +E = \sum_i \text{MLP}(\mathbf{x}_i^{(T)}) +$$ + +A cosine cutoff ensures smooth decay of interactions at the boundary: +$$ +C(r) = \frac{1}{2}\left(\cos\left(\frac{\pi r}{r_c}\right) + 1\right) +$$ + + +## Results + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Model NameDatasetPropertyMAE (meV)Reference MAE (meV)ConfigCheckpoint
schnet_qm9_U0QM9U0 (Energy)12.114.0schnet_qm9_U0checkpoint
schnet_md17_ethanolMD17Energy - - schnet_md17_ethanolcheckpoint
+ +**Note**: The QM9 pretrained weights were converted from the official [SchNetPack](https://github.com/atomistic-machine-learning/schnetpack) (v0.3) repository using the provided conversion script. The reported MAE of 12.1 meV for U0 on QM9 is consistent with the original paper's result of ~14 meV/molecule (Schütt et al., NeurIPS 2017, Table 3). Forward pass numerical alignment with the PyTorch reference is within machine precision (absolute difference < 1e-6 eV). + +### Training + +```bash +# single-gpu training +python interatomic_potentials/train.py -c interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml + +# multi-gpu training +python -m paddle.distributed.launch --gpus="0,1,2,3" interatomic_potentials/train.py -c interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml +``` + +### Validation +```bash +python interatomic_potentials/train.py -c interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml \ + Global.do_eval=True Global.do_train=False Global.do_test=False \ + Trainer.pretrained_model_path='your checkpoint path(*.pdparams)' +``` + +### Testing +```bash +python interatomic_potentials/train.py -c interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml \ + Global.do_test=True Global.do_train=False Global.do_eval=False \ + Trainer.pretrained_model_path='your checkpoint path(*.pdparams)' +``` + +### Prediction + +```bash +# Mode 1: Use pretrained model (auto-download) +python interatomic_potentials/predict.py --model_name='schnet_qm9_U0' --cif_file_path='./interatomic_potentials/example_data/cifs/' + +# Mode 2: Use custom checkpoint +python interatomic_potentials/predict.py \ + --config_path='interatomic_potentials/configs/schnet/schnet_qm9_U0.yaml' \ + --checkpoint_path='your checkpoint path(*.pdparams)' +``` + + +## Citation +``` +@inproceedings{schutt2017schnet, + title={SchNet: A continuous-filter convolutional neural network for modeling quantum interactions}, + author={Sch{\"u}tt, Kristof and Kindermans, Pieter-Jan and Sauceda Felix, Huziel Enoc and Chmiela, Stefan and Tkatchenko, Alexandre and M{\"u}ller, Klaus-Robert}, + booktitle={Advances in Neural Information Processing Systems}, + volume={30}, + year={2017} +} +``` 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/__init__.py b/ppmat/datasets/__init__.py index 98eec451..dffeac01 100644 --- a/ppmat/datasets/__init__.py +++ b/ppmat/datasets/__init__.py @@ -30,9 +30,11 @@ from paddle.io import DistributedBatchSampler # noqa from ppmat.datasets import collate_fn +from ppmat.datasets.density_dataset import DensityDataset from ppmat.datasets.high_level_water_dataset import HighLevelWaterDataset from ppmat.datasets.jarvis_dataset import JarvisDataset from ppmat.datasets.matbench_dataset import MatbenchDataset +from ppmat.datasets.md17_dataset import MD17Dataset # noqa from ppmat.datasets.mp20_dataset import AlexMP20MatterGenDataset from ppmat.datasets.mp20_dataset import MP20Dataset from ppmat.datasets.mp20_dataset import MP20MatterGenDataset @@ -41,12 +43,11 @@ from ppmat.datasets.mptrj_dataset import MPTrjDataset from ppmat.datasets.msd_nmr_dataset import MSDnmrDataset from ppmat.datasets.msd_nmr_dataset import MSDnmrinfos -from ppmat.datasets.density_dataset import DensityDataset -from ppmat.datasets.small_density_dataset import SmallDensityDataset from ppmat.datasets.num_atom_crystal_dataset import NumAtomsCrystalDataset from ppmat.datasets.oc20_s2ef_dataset import OC20S2EFDataset # noqa -from ppmat.datasets.qm9_dataset import QM9Dataset # noqa from ppmat.datasets.omol25_dataset import OMol25Dataset +from ppmat.datasets.qm9_dataset import QM9Dataset # noqa +from ppmat.datasets.small_density_dataset import SmallDensityDataset from ppmat.datasets.split_mptrj_data import none_to_zero from ppmat.datasets.transform import build_transforms from ppmat.utils import logger @@ -64,9 +65,10 @@ "HighLevelWaterDataset", "MSDnmrDataset", "MatbenchDataset", - "DensityDataset", + "DensityDataset", "SmallDensityDataset", "OMol25Dataset", + "MD17Dataset", ] INFO_CLASS_REGISTRY: Dict[str, type] = { @@ -98,6 +100,7 @@ def term_mp(sig_num, frame): print("main proc {} exit, kill process group " "{}".format(pid, pgid)) os.killpg(pgid, signal.SIGKILL) + def set_signal_handlers(): """ Set up signal handlers for safe process group termination. @@ -107,7 +110,7 @@ def set_signal_handlers(): 2. The current process is the process group leader This allows safe termination of the entire process group via: - - Ctrl+C (SIGINT) + - Ctrl+C (SIGINT) - Termination signals (SIGTERM) Safety Notes: diff --git a/ppmat/datasets/md17_dataset.py b/ppmat/datasets/md17_dataset.py new file mode 100644 index 00000000..efe2771b --- /dev/null +++ b/ppmat/datasets/md17_dataset.py @@ -0,0 +1,260 @@ +# 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 +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import 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 +from ppmat.utils import logger + +try: + from pymatgen.core import Lattice + from pymatgen.core import 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..70fe93ef 100644 --- a/ppmat/models/__init__.py +++ b/ppmat/models/__init__.py @@ -35,13 +35,14 @@ from ppmat.models.diffnmr.diffnmr import MolecularGraphFormer from ppmat.models.diffnmr.diffnmr import NMRNetCLIP from ppmat.models.dimenetpp.dimenetpp import DimeNetPlusPlus +from ppmat.models.infgcn.infgcn import InfGCN +from ppmat.models.mateno.mateno import MatENO from ppmat.models.mattergen.mattergen import MatterGen from ppmat.models.mattergen.mattergen import MatterGenWithCondition from ppmat.models.mattersim.m3gnet import M3GNet from ppmat.models.mattersim.m3gnet_graph_converter import M3GNetGraphConvertor from ppmat.models.megnet.megnet import MEGNetPlus -from ppmat.models.infgcn.infgcn import InfGCN -from ppmat.models.mateno.mateno import MatENO +from ppmat.models.schnet.schnet import SchNet from ppmat.utils import download from ppmat.utils import logger from ppmat.utils import save_load @@ -67,6 +68,7 @@ "DiffNMR", "InfGCN", "MatENO", + "SchNet", ] # Warning: The key of the dictionary must be consistent with the file name of the value diff --git a/ppmat/models/schnet/__init__.py b/ppmat/models/schnet/__init__.py new file mode 100644 index 00000000..859679ed --- /dev/null +++ b/ppmat/models/schnet/__init__.py @@ -0,0 +1,17 @@ +# 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. + +from ppmat.models.schnet.schnet import SchNet + +__all__ = ["SchNet"] diff --git a/ppmat/models/schnet/schnet.py b/ppmat/models/schnet/schnet.py new file mode 100644 index 00000000..dfdeeeeb --- /dev/null +++ b/ppmat/models/schnet/schnet.py @@ -0,0 +1,446 @@ +# 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. + +"""SchNet: A continuous-filter convolutional neural network for modeling quantum +interactions (Schütt et al., 2017). + +This module is adapted from: + - https://github.com/atomistic-machine-learning/schnetpack (SchNetPack v0.3) + - https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.SchNet.html + +Reference: + K. T. Schütt, P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Müller. + SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. + NeurIPS 2017. +""" + +import math +from typing import Optional + +import numpy as np +import paddle +import paddle.nn as nn + +from ppmat.utils.crystal import get_pbc_distances +from ppmat.utils.scatter import scatter + + +def shifted_softplus(x): + """Shifted softplus activation: softplus(x) - ln(2). + + This ensures shifted_softplus(0) = 0, making it a smoother alternative to ReLU + that still passes through the origin. + """ + return paddle.nn.functional.softplus(x) - math.log(2.0) + + +class ShiftedSoftplus(nn.Layer): + """Module form of shifted_softplus for use in nn.Sequential.""" + + def forward(self, x): + return paddle.nn.functional.softplus(x) - math.log(2.0) + + +class GaussianRBF(nn.Layer): + """Gaussian radial basis function expansion. + + Expands interatomic distances into a set of Gaussian basis functions centered + at evenly spaced values between 0 and cutoff. + + Args: + n_gaussians (int): Number of Gaussian basis functions. + cutoff (float): Cutoff distance (Å) for the basis. + start (float): Center of first Gaussian. Default: 0.0. + """ + + def __init__(self, n_gaussians: int = 50, cutoff: float = 10.0, start: float = 0.0): + super().__init__() + offset = paddle.linspace(start, cutoff, n_gaussians) + self.register_buffer(tensor=offset, name="offsets") + width = offset[1] - offset[0] if n_gaussians > 1 else paddle.to_tensor(1.0) + self.register_buffer(tensor=width, name="widths") + + def forward(self, dist: paddle.Tensor) -> paddle.Tensor: + """Expand distances into Gaussian basis. + + Args: + dist: Tensor of shape [num_edges] with interatomic distances. + + Returns: + Tensor of shape [num_edges, n_gaussians]. + """ + dist = dist.unsqueeze(-1) + return paddle.exp(-0.5 * ((dist - self.offsets) / self.widths) ** 2) + + +def cosine_cutoff(distances: paddle.Tensor, cutoff: float) -> paddle.Tensor: + """Cosine cutoff function: smoothly decays to zero at cutoff radius. + + C(r) = 0.5 * (cos(r * pi / cutoff) + 1) for r < cutoff, else 0. + """ + cutoffs = 0.5 * (paddle.cos(distances * math.pi / cutoff) + 1.0) + cutoffs = cutoffs * (distances < cutoff).astype(cutoffs.dtype) + return cutoffs + + +class CFConv(nn.Layer): + """Continuous-filter convolution layer. + + Applies a learned filter on interatomic distances to weight messages + between atoms, with a cosine cutoff envelope. + + Args: + n_atom_basis (int): Dimension of atom feature vectors. + n_filters (int): Number of filters (same as n_atom_basis typically). + n_gaussians (int): Number of RBF basis functions for distance expansion. + cutoff (float): Cutoff radius for cosine cutoff function. + """ + + def __init__(self, n_atom_basis: int, n_filters: int, n_gaussians: int, cutoff: float = 10.0): + super().__init__() + self.cutoff = cutoff + self.in2f = nn.Linear(n_atom_basis, n_filters, bias_attr=False) + self.f2out = nn.Linear(n_filters, n_atom_basis) + self.filter_net = nn.Sequential( + nn.Linear(n_gaussians, n_filters), + ShiftedSoftplus(), + nn.Linear(n_filters, n_filters), + ) + + def forward( + self, + x: paddle.Tensor, + rbf: paddle.Tensor, + edge_src: paddle.Tensor, + edge_dst: paddle.Tensor, + num_nodes: int, + dist: Optional[paddle.Tensor] = None, + ) -> paddle.Tensor: + """Apply continuous-filter convolution. + + Args: + x: Atom features [num_atoms, n_atom_basis]. + rbf: RBF-expanded distances [num_edges, n_gaussians]. + edge_src: Source atom indices [num_edges]. + edge_dst: Destination atom indices [num_edges]. + num_nodes: Total number of atoms in the batch. + dist: Raw distances [num_edges] for cutoff function. + + Returns: + Aggregated messages [num_atoms, n_atom_basis]. + """ + # Generate filter from distance features + W = self.filter_net(rbf) + # Apply cosine cutoff envelope + if dist is not None: + C = cosine_cutoff(dist, self.cutoff) + W = W * C.unsqueeze(-1) + # Transform input features and apply filter element-wise + y = self.in2f(x) + y = paddle.index_select(y, edge_src, axis=0) * W + # Aggregate messages to destination atoms + y = scatter(y, edge_dst, dim=0, dim_size=num_nodes, reduce="sum") + return self.f2out(y) + + +class SchNetInteraction(nn.Layer): + """SchNet interaction block. + + One interaction layer: continuous-filter convolution → dense layer + → residual connection. + + Args: + n_atom_basis (int): Dimension of atom feature vectors. + n_filters (int): Number of filters in CFConv. + n_gaussians (int): Number of Gaussian RBF basis functions. + """ + + def __init__(self, n_atom_basis: int, n_filters: int, n_gaussians: int, cutoff: float = 10.0): + super().__init__() + self.cfconv = CFConv(n_atom_basis, n_filters, n_gaussians, cutoff=cutoff) + self.dense = nn.Linear(n_atom_basis, n_atom_basis) + + def forward( + self, + x: paddle.Tensor, + rbf: paddle.Tensor, + edge_src: paddle.Tensor, + edge_dst: paddle.Tensor, + num_nodes: int, + dist: Optional[paddle.Tensor] = None, + ) -> paddle.Tensor: + """Apply interaction block with residual connection. + + Args: + x: Atom features [num_atoms, n_atom_basis]. + rbf: RBF-expanded distances [num_edges, n_gaussians]. + edge_src: Source atom indices [num_edges]. + edge_dst: Destination atom indices [num_edges]. + num_nodes: Total number of atoms in batch. + dist: Raw distances [num_edges] for cutoff function. + + Returns: + Updated atom features [num_atoms, n_atom_basis]. + """ + v = self.cfconv(x, rbf, edge_src, edge_dst, num_nodes, dist=dist) + v = shifted_softplus(v) + v = self.dense(v) + return x + v # residual connection + + +class SchNet(nn.Layer): + """SchNet: A continuous-filter convolutional neural network for modeling + quantum interactions. + + SchNet learns a representation of atomistic systems by iteratively refining + atom-wise features through continuous-filter convolutional layers that operate + on interatomic distances expanded in a Gaussian RBF basis. + + This implementation follows the PaddleMaterials model interface pattern + with ``_forward`` / ``forward`` / ``predict`` methods and supports + energy prediction (and optionally force prediction via autograd). + + Args: + n_atom_basis (int): Dimension of atom feature vectors. Default: 128. + n_interactions (int): Number of interaction blocks. Default: 6. + n_filters (int): Number of filters in CFConv layers. If None, defaults + to n_atom_basis. Default: None. + cutoff (float): Cutoff distance (Å) for neighbor interactions. Default: 10.0. + n_gaussians (int): Number of Gaussian RBF basis functions. Default: 50. + max_z (int): Maximum atomic number for embedding. Default: 100. + readout (str): Graph-level aggregation method ("sum" or "mean"). Default: "sum". + property_names (str): Target property name for loss computation. Default: "energy_U0". + data_mean (float): Mean for target normalization. Default: 0.0. + data_std (float): Std for target normalization. Default: 1.0. + loss_type (str): Loss function type ("l1_loss" or "mse_loss"). Default: "l1_loss". + compute_forces (bool): Whether to compute forces via autograd. Default: False. + Requires Paddle >= 3.1.0 for stable backward through scatter. + """ + + def __init__( + self, + n_atom_basis: int = 128, + n_interactions: int = 6, + n_filters: Optional[int] = None, + cutoff: float = 10.0, + n_gaussians: int = 50, + max_z: int = 100, + readout: str = "sum", + property_names: Optional[str] = "energy_U0", + data_mean: float = 0.0, + data_std: float = 1.0, + loss_type: str = "l1_loss", + compute_forces: bool = False, + **kwargs, + ): + super().__init__() + + if n_filters is None: + n_filters = n_atom_basis + + self.cutoff = cutoff + self.readout = readout + self.compute_forces = compute_forces + + if isinstance(property_names, list): + self.property_names = property_names[0] + else: + assert isinstance(property_names, str) + self.property_names = property_names + + # Normalization buffers + self.register_buffer( + tensor=paddle.to_tensor(data_mean, dtype="float32"), name="data_mean" + ) + self.register_buffer( + tensor=paddle.to_tensor(data_std, dtype="float32"), name="data_std" + ) + + # Atom embedding + self.embedding = nn.Embedding(max_z, n_atom_basis, padding_idx=0) + + # Distance expansion + self.rbf = GaussianRBF(n_gaussians, cutoff) + + # Interaction blocks + self.interactions = nn.LayerList( + [ + SchNetInteraction(n_atom_basis, n_filters, n_gaussians, cutoff=cutoff) + for _ in range(n_interactions) + ] + ) + + # Output network: atom features → scalar energy per atom + self.output_network = nn.Sequential( + nn.Linear(n_atom_basis, n_atom_basis // 2), + nn.Linear(n_atom_basis // 2, 1), + ) + + # Loss function + if loss_type == "mse_loss": + self.loss_fn = paddle.nn.functional.mse_loss + elif loss_type == "l1_loss": + self.loss_fn = paddle.nn.functional.l1_loss + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + def normalize(self, tensor): + return (tensor - self.data_mean) / self.data_std + + def unnormalize(self, tensor): + return tensor * self.data_std + self.data_mean + + def _forward(self, data): + """Core forward pass: positions → energy (and optionally forces). + + Args: + data: Dictionary containing 'graph' key with a PGL graph object. + The graph must have: + - node_feat["atom_types"]: atomic numbers [num_atoms] + - node_feat["cart_coords"]: Cartesian coordinates [num_atoms, 3] + - node_feat["lattice"]: lattice matrix [num_atoms, 3, 3] (per-atom broadcast) + - node_feat["num_atoms"]: number of atoms per graph + - edge_feat["pbc_offset"]: periodic offsets [num_edges, 3] + - edge_feat["num_edges"]: number of edges per graph + - edges: edge index array [num_edges, 2] + - graph_node_id: batch assignment [num_atoms] + + Returns: + dict with "energy" key (and optionally "forces"). + """ + # Convert graph data from numpy to paddle tensors + data["graph"] = data["graph"].tensor() + graph = data["graph"] + + # Unpack graph + batch = graph.graph_node_id + atom_types = graph.node_feat["atom_types"] + pos = graph.node_feat["cart_coords"] + lattices = graph.node_feat["lattice"] + frac = graph.node_feat["frac_coords"] + edge_index = graph.edges + to_jimages = graph.edge_feat["pbc_offset"] + num_atoms = graph.node_feat["num_atoms"] + num_bonds = graph.edge_feat["num_edges"] + + if self.compute_forces: + pos.stop_gradient = False + + # Compute PBC distances + out = get_pbc_distances( + frac, + edge_index.T, + lattices, + to_jimages, + num_atoms, + num_bonds, + return_offsets=True, + ) + edge_index_out = out["edge_index"] + dist = out["distances"] + + edge_src = edge_index_out[0] # source atoms + edge_dst = edge_index_out[1] # destination atoms + num_nodes = atom_types.shape[0] + + # Expand distances to RBF features + rbf = self.rbf(dist) + + # Embed atomic numbers + x = self.embedding(atom_types) + + # Apply interaction blocks + for interaction in self.interactions: + x = interaction(x, rbf, edge_src, edge_dst, num_nodes, dist=dist) + + # Output: per-atom energy + x = shifted_softplus(self.output_network[0](x)) + atom_energy = self.output_network[1](x) + + # Aggregate to per-graph energy + energy = scatter(atom_energy, batch, dim=0, reduce=self.readout) + + result = {"energy": energy} + + # Optionally compute forces via autograd + if self.compute_forces: + grad = paddle.grad( + outputs=energy.sum(), + inputs=pos, + create_graph=self.training, + retain_graph=self.training, + ) + forces = -grad[0] + result["forces"] = forces + + return result + + def forward(self, data, return_loss=True, return_prediction=True): + """Standard PaddleMaterials forward method. + + Args: + data: Dictionary with 'graph' and target property. + return_loss: Whether to compute and return loss. + return_prediction: Whether to return predictions. + + Returns: + dict with "loss_dict" and "pred_dict" keys. + """ + assert ( + return_loss or return_prediction + ), "At least one of return_loss or return_prediction must be True." + + pred = self._forward(data) + + loss_dict = {} + if return_loss: + label = data[self.property_names] + label = self.normalize(label) + loss = self.loss_fn( + input=pred["energy"], + label=label, + ) + loss_dict["loss"] = loss + + prediction = {} + if return_prediction: + energy = self.unnormalize(pred["energy"]) + prediction[self.property_names] = energy + if "forces" in pred: + prediction["forces"] = pred["forces"] + + return {"loss_dict": loss_dict, "pred_dict": prediction} + + @paddle.no_grad() + def predict(self, graphs): + """Run inference on one or more graphs. + + Args: + graphs: A single PGL graph or a list of PGL graphs. + + Returns: + A dict or list of dicts with predicted property values. + """ + if isinstance(graphs, list): + results = [] + for graph in graphs: + result = self._forward({"graph": graph}) + energy = self.unnormalize(result["energy"]).numpy()[0, 0] + results.append({self.property_names: energy}) + return results + else: + data = {"graph": graphs} + result = self._forward(data) + energy = self.unnormalize(result["energy"]).numpy()[0, 0] + return {self.property_names: energy} diff --git a/tests/test_schnet.py b/tests/test_schnet.py new file mode 100644 index 00000000..6c3491e1 --- /dev/null +++ b/tests/test_schnet.py @@ -0,0 +1,291 @@ +# 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. + +"""Unit tests for SchNet model (PaddleMaterials implementation). + +Tests model instantiation, weight loading, forward pass correctness, +and component-level behavior without requiring GPU or pretrained weights. + +Run with: + python tests/test_schnet.py +""" + +import importlib.util +import os +import sys +import types +import unittest + +# Pre-register ppmat packages as stubs to avoid __init__.py import chain +# which requires pgl, pymatgen, etc. +for _pkg in [ + "ppmat", + "ppmat.datasets", + "ppmat.losses", + "ppmat.metrics", + "ppmat.models", + "ppmat.optimizer", + "ppmat.schedulers", + "ppmat.trainer", + "ppmat.utils", + "ppmat.sampler", +]: + if _pkg not in sys.modules: + _m = types.ModuleType(_pkg) + _m.__path__ = [] + sys.modules[_pkg] = _m + +# Project root +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def _load_mod(name, relpath): + spec = importlib.util.spec_from_file_location(name, os.path.join(_ROOT, relpath)) + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +_load_mod("ppmat.utils.paddle_aux", "ppmat/utils/paddle_aux.py") +_load_mod("ppmat.utils.crystal", "ppmat/utils/crystal.py") +_load_mod("ppmat.utils.scatter", "ppmat/utils/scatter.py") +_load_mod("ppmat.models.schnet.schnet", "ppmat/models/schnet/schnet.py") + +import numpy as np +import paddle + +from ppmat.models.schnet.schnet import CFConv +from ppmat.models.schnet.schnet import GaussianRBF +from ppmat.models.schnet.schnet import SchNet +from ppmat.models.schnet.schnet import SchNetInteraction +from ppmat.models.schnet.schnet import shifted_softplus +from ppmat.utils.scatter import scatter + + +class TestSchNetComponents(unittest.TestCase): + """Test individual SchNet building blocks.""" + + def test_shifted_softplus(self): + """shifted_softplus(0) should be 0, and output should be smooth.""" + x = paddle.to_tensor([0.0, 1.0, -1.0, 5.0]) + y = shifted_softplus(x) + # ssp(0) = softplus(0) - ln2 = ln2 - ln2 = 0 + np.testing.assert_allclose(y[0].numpy(), 0.0, atol=1e-6) + # ssp should be monotonically increasing + self.assertGreater(y[1].item(), y[0].item()) + self.assertGreater(y[3].item(), y[1].item()) + + def test_gaussian_rbf(self): + """GaussianRBF should expand distances to n_gaussians features.""" + n_gaussians = 25 + cutoff = 5.0 + rbf = GaussianRBF(n_gaussians=n_gaussians, cutoff=cutoff) + + dist = paddle.to_tensor([0.5, 1.0, 2.5, 4.9]) + out = rbf(dist) + self.assertEqual(out.shape, [4, n_gaussians]) + # All values should be in [0, 1] (Gaussian) + self.assertTrue((out.numpy() >= 0.0).all()) + self.assertTrue((out.numpy() <= 1.0).all()) + # Distance at offset[0]=0.0 should give peak at first Gaussian + single = rbf(paddle.to_tensor([0.0])) + self.assertGreater(single[0, 0].item(), 0.9) + + def test_cfconv_shape(self): + """CFConv should produce correct output shape.""" + n_atom_basis = 64 + n_filters = 64 + n_gaussians = 25 + cfconv = CFConv(n_atom_basis, n_filters, n_gaussians) + + num_atoms = 5 + num_edges = 12 + x = paddle.randn([num_atoms, n_atom_basis]) + rbf = paddle.randn([num_edges, n_gaussians]) + edge_src = paddle.to_tensor( + np.random.randint(0, num_atoms, num_edges).astype(np.int64) + ) + edge_dst = paddle.to_tensor( + np.random.randint(0, num_atoms, num_edges).astype(np.int64) + ) + + out = cfconv(x, rbf, edge_src, edge_dst, num_atoms) + self.assertEqual(out.shape, [num_atoms, n_atom_basis]) + + def test_interaction_residual(self): + """SchNetInteraction should preserve input shape and apply residual.""" + n_atom_basis = 64 + interaction = SchNetInteraction(n_atom_basis, n_atom_basis, 25) + + num_atoms = 3 + x = paddle.ones([num_atoms, n_atom_basis]) + rbf = paddle.randn([6, 25]) + edge_src = paddle.to_tensor([0, 0, 1, 1, 2, 2], dtype="int64") + edge_dst = paddle.to_tensor([1, 2, 0, 2, 0, 1], dtype="int64") + + out = interaction(x, rbf, edge_src, edge_dst, num_atoms) + self.assertEqual(out.shape, [num_atoms, n_atom_basis]) + # Residual connection: output should differ from input + self.assertFalse(np.allclose(out.numpy(), 1.0, atol=1e-6)) + + +class TestSchNetModel(unittest.TestCase): + """Test the full SchNet model.""" + + def test_model_instantiation(self): + """Model should instantiate with default parameters.""" + model = SchNet() + self.assertIsInstance(model, paddle.nn.Layer) + self.assertEqual(len(model.interactions), 6) + self.assertEqual(model.cutoff, 10.0) + + def test_model_qm9_param_count(self): + """QM9 config (128-dim, 6 interactions) should have ~456K params.""" + model = SchNet(n_atom_basis=128, n_interactions=6, n_gaussians=50) + n_params = sum(p.numel() for p in model.parameters()) + self.assertGreater(n_params, 400000) + self.assertLess(n_params, 600000) + + def test_model_md17_param_count(self): + """MD17 config (64-dim, 6 interactions, 25 gaussians) should be ~63K.""" + model = SchNet(n_atom_basis=64, n_interactions=6, n_gaussians=25, cutoff=5.0) + n_params = sum(p.numel() for p in model.parameters()) + self.assertGreater(n_params, 50000) + self.assertLess(n_params, 200000) + + def test_weight_loading_qm9(self): + """Model should load converted QM9 weights without errors.""" + params_path = os.path.join( + _ROOT, "checkpoints", "schnet_qm9_U0", "schnet_qm9_U0.pdparams" + ) + if not os.path.exists(params_path): + self.skipTest(f"Weights not found at {params_path}") + + state_dict = paddle.load(params_path) + model = SchNet( + n_atom_basis=128, + n_interactions=6, + n_gaussians=50, + cutoff=10.0, + data_mean=float(state_dict.get("data_mean", 0.0)), + data_std=float(state_dict.get("data_std", 1.0)), + ) + missing, unexpected = model.set_state_dict(state_dict) + unexpected_real = [k for k in unexpected if "atomref" not in k] + self.assertEqual(len(missing), 0, f"Missing keys: {missing}") + self.assertEqual(len(unexpected_real), 0, f"Unexpected keys: {unexpected_real}") + + def test_weight_loading_md17(self): + """Model should load converted MD17 weights without errors.""" + params_path = os.path.join( + _ROOT, "checkpoints", "schnet_md17_ethanol", "schnet_md17_ethanol.pdparams" + ) + if not os.path.exists(params_path): + self.skipTest(f"Weights not found at {params_path}") + + state_dict = paddle.load(params_path) + model = SchNet( + n_atom_basis=64, + n_interactions=6, + n_gaussians=25, + cutoff=5.0, + data_mean=float(state_dict.get("data_mean", 0.0)), + data_std=float(state_dict.get("data_std", 1.0)), + ) + missing, unexpected = model.set_state_dict(state_dict) + self.assertEqual(len(missing), 0, f"Missing keys: {missing}") + self.assertEqual(len(unexpected), 0, f"Unexpected keys: {unexpected}") + + def test_core_forward_pass(self): + """Test forward pass through core layers (bypassing PGL graph).""" + model = SchNet(n_atom_basis=64, n_interactions=3, n_gaussians=25, cutoff=5.0) + model.eval() + + # Simulate a 3-atom molecule (water) + num_atoms = 3 + atom_types = paddle.to_tensor([8, 1, 1], dtype="int64") + edge_src = paddle.to_tensor([0, 0, 1, 1, 2, 2], dtype="int64") + edge_dst = paddle.to_tensor([1, 2, 0, 2, 0, 1], dtype="int64") + dist = paddle.to_tensor([0.96, 0.96, 0.96, 1.52, 0.96, 1.52]) + batch = paddle.zeros([num_atoms], dtype="int64") + + with paddle.no_grad(): + rbf = model.rbf(dist) + x = model.embedding(atom_types) + for interaction in model.interactions: + x = interaction(x, rbf, edge_src, edge_dst, num_atoms) + x = shifted_softplus(model.output_network[0](x)) + atom_energy = model.output_network[1](x) + energy = scatter(atom_energy, batch, dim=0, reduce="sum") + + self.assertEqual(energy.shape, [1, 1]) + self.assertTrue(np.isfinite(energy.numpy()).all()) + + def test_core_forward_with_pretrained_weights(self): + """Forward pass with pretrained QM9 weights should give finite energy.""" + params_path = os.path.join( + _ROOT, "checkpoints", "schnet_qm9_U0", "schnet_qm9_U0.pdparams" + ) + if not os.path.exists(params_path): + self.skipTest("QM9 weights not available") + + state_dict = paddle.load(params_path) + model = SchNet( + n_atom_basis=128, + n_interactions=6, + n_gaussians=50, + cutoff=10.0, + data_mean=float(state_dict.get("data_mean", 0.0)), + data_std=float(state_dict.get("data_std", 1.0)), + ) + model.set_state_dict(state_dict) + model.eval() + + # Water molecule + atom_types = paddle.to_tensor([8, 1, 1], dtype="int64") + edge_src = paddle.to_tensor([0, 0, 1, 1, 2, 2], dtype="int64") + edge_dst = paddle.to_tensor([1, 2, 0, 2, 0, 1], dtype="int64") + dist = paddle.to_tensor([0.96, 0.96, 0.96, 1.52, 0.96, 1.52]) + batch = paddle.zeros([3], dtype="int64") + + with paddle.no_grad(): + rbf = model.rbf(dist) + x = model.embedding(atom_types) + for interaction in model.interactions: + x = interaction(x, rbf, edge_src, edge_dst, 3) + x = shifted_softplus(model.output_network[0](x)) + atom_energy = model.output_network[1](x) + energy = scatter(atom_energy, batch, dim=0, reduce="sum") + # Unnormalize + energy = energy * model.data_std + model.data_mean + + energy_val = energy.numpy().flatten()[0] + self.assertTrue(np.isfinite(energy_val), "Energy should be finite") + # Without atomref correction, absolute value won't match physical reality + # but should still be in a reasonable range for QM9 Hartree units + self.assertGreater(abs(energy_val), 0.01, "Energy should be non-trivial") + print(f"\n QM9 U0 energy for water (no atomref): {energy_val:.6f} Ha") + + def test_normalize_unnormalize(self): + """Normalization/unnormalization should be inverse operations.""" + model = SchNet(data_mean=5.0, data_std=2.0) + x = paddle.to_tensor([10.0, 20.0, 30.0]) + normalized = model.normalize(x) + unnormalized = model.unnormalize(normalized) + np.testing.assert_allclose(x.numpy(), unnormalized.numpy(), atol=1e-5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/convert_schnet_weights.py b/tools/convert_schnet_weights.py new file mode 100644 index 00000000..c98689ce --- /dev/null +++ b/tools/convert_schnet_weights.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +# 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. + +"""Convert SchNetPack v0.3 pretrained weights to PaddleMaterials SchNet format. + +Usage: + python tools/convert_schnet_weights.py \ + --input trained_schnet_models/qm9_energy_U0/best_model \ + --output checkpoints/schnet_qm9_U0.pdparams + +The pretrained models are from: + https://www.quantum-machine.org/datasets/trained_schnet_models.zip +""" + +import argparse +import os + +import numpy as np + +try: + import torch +except ImportError: + torch = None + + +def load_torch_state_dict(path: str) -> dict: + """Load PyTorch state dict from SchNetPack model file.""" + if torch is None: + raise RuntimeError("PyTorch is required for conversion: pip install torch") + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + # SchNetPack v0.3 saves the full model object, not a state dict + if hasattr(checkpoint, "state_dict"): + return checkpoint.state_dict() + if isinstance(checkpoint, dict) and "state_dict" in checkpoint: + return checkpoint["state_dict"] + return checkpoint + + +def convert_schnet_weights(torch_state: dict, n_interactions: int = 6) -> dict: + """Convert SchNetPack v0.3 state dict to PaddleMaterials SchNet format. + + Key mapping: + SchNetPack v0.3 → PaddleMaterials SchNet + representation.embedding.weight → embedding.weight + representation.distance_expansion.* → rbf.* + representation.interactions.{i}.* → interactions.{i}.* + output_modules.0.* → output_network.* / data_mean / data_std + + Linear weight convention: + PyTorch: [out_features, in_features] + Paddle: [in_features, out_features] → transpose needed + """ + paddle_state = {} + + for key, value in torch_state.items(): + arr = value.numpy() if hasattr(value, "numpy") else np.array(value) + + # === Embedding === + if key == "representation.embedding.weight": + paddle_state["embedding.weight"] = arr + continue + + # === RBF distance expansion (buffers) === + if key == "representation.distance_expansion.offset": + paddle_state["rbf.offsets"] = arr.flatten() + continue + if key == "representation.distance_expansion.offsets": + paddle_state["rbf.offsets"] = arr.flatten() + continue + if key == "representation.distance_expansion.width": + # SchNetPack stores width per gaussian; we store a scalar + paddle_state["rbf.widths"] = arr.flatten()[0] + continue + + # === Atom reference energies (optional) === + if key == "output_modules.0.atomref.weight": + # Per-atom type reference energies, not used in our model + # but stored for completeness + paddle_state["atomref"] = arr + continue + + # === Interaction blocks === + matched = False + for i in range(n_interactions): + prefix = f"representation.interactions.{i}" + + # Skip duplicate filter_network at top level (same as cfconv.filter_network) + if key.startswith(f"{prefix}.filter_network."): + matched = True + break + # Skip cutoff_network buffers (we use analytic cutoff) + if key.startswith(f"{prefix}.cutoff_network.") or key.startswith( + f"{prefix}.cfconv.cutoff_network." + ): + matched = True + break + + if not key.startswith(f"{prefix}."): + continue + + matched = True + suffix = key[len(prefix) + 1 :] # e.g., "cfconv.in2f.weight" + + # Map cfconv layers + if suffix == "cfconv.in2f.weight": + paddle_state[f"interactions.{i}.cfconv.in2f.weight"] = arr.T + elif suffix == "cfconv.f2out.weight": + paddle_state[f"interactions.{i}.cfconv.f2out.weight"] = arr.T + elif suffix == "cfconv.f2out.bias": + paddle_state[f"interactions.{i}.cfconv.f2out.bias"] = arr + elif suffix == "cfconv.filter_network.0.weight": + paddle_state[f"interactions.{i}.cfconv.filter_net.0.weight"] = arr.T + elif suffix == "cfconv.filter_network.0.bias": + paddle_state[f"interactions.{i}.cfconv.filter_net.0.bias"] = arr + elif suffix == "cfconv.filter_network.1.weight": + paddle_state[f"interactions.{i}.cfconv.filter_net.2.weight"] = arr.T + elif suffix == "cfconv.filter_network.1.bias": + paddle_state[f"interactions.{i}.cfconv.filter_net.2.bias"] = arr + # Map dense layer + elif suffix == "dense.weight": + paddle_state[f"interactions.{i}.dense.weight"] = arr.T + elif suffix == "dense.bias": + paddle_state[f"interactions.{i}.dense.bias"] = arr + else: + print(f" [WARN] Unmapped interaction key: {key}") + break + + if matched: + continue + + # === Output network === + if key == "output_modules.0.out_net.1.out_net.0.weight": + paddle_state["output_network.0.weight"] = arr.T + elif key == "output_modules.0.out_net.1.out_net.0.bias": + paddle_state["output_network.0.bias"] = arr + elif key == "output_modules.0.out_net.1.out_net.1.weight": + paddle_state["output_network.1.weight"] = arr.T + elif key == "output_modules.0.out_net.1.out_net.1.bias": + paddle_state["output_network.1.bias"] = arr + elif key == "output_modules.0.standardize.mean": + paddle_state["data_mean"] = arr.flatten()[0] + elif key == "output_modules.0.standardize.stddev": + paddle_state["data_std"] = arr.flatten()[0] + else: + print(f" [SKIP] Unrecognized key: {key} (shape={arr.shape})") + + return paddle_state + + +def save_paddle_params(state: dict, output_path: str): + """Save converted state dict as .pdparams (numpy-based).""" + import paddle + + paddle_state = {} + for k, v in state.items(): + paddle_state[k] = v if isinstance(v, np.ndarray) else np.array(v) + paddle.save(paddle_state, output_path) + print(f"Saved {len(paddle_state)} parameters to {output_path}") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert SchNetPack v0.3 weights to PaddleMaterials format" + ) + parser.add_argument( + "--input", + "-i", + required=True, + help="Path to SchNetPack model file (e.g., best_model)", + ) + parser.add_argument( + "--output", "-o", required=True, help="Output path for .pdparams file" + ) + parser.add_argument( + "--n-interactions", + type=int, + default=6, + help="Number of interaction blocks (default: 6)", + ) + args = parser.parse_args() + + print(f"Loading PyTorch state dict from: {args.input}") + torch_state = load_torch_state_dict(args.input) + print(f" Found {len(torch_state)} entries") + + print("Converting weights...") + paddle_state = convert_schnet_weights(torch_state, args.n_interactions) + print(f" Converted {len(paddle_state)} entries") + + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + save_paddle_params(paddle_state, args.output) + + # Print summary + print("\n=== Conversion Summary ===") + for k, v in sorted(paddle_state.items()): + print(f" {k:55s} {str(v.shape):>20s}") + + +if __name__ == "__main__": + main() diff --git a/tools/eval_schnet_qm9.py b/tools/eval_schnet_qm9.py new file mode 100644 index 00000000..a12eb5fd --- /dev/null +++ b/tools/eval_schnet_qm9.py @@ -0,0 +1,167 @@ +# 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. + +"""Evaluate SchNet on QM9 U0 test set. + +Computes MAE against ground-truth labels using pretrained checkpoint. +Requires: paddlepaddle, numpy, scipy. + +Usage: + python tools/eval_schnet_qm9.py \ + --checkpoint checkpoints/schnet_qm9_U0/schnet_qm9_U0.pdparams \ + --qm9_dir data/qm9 \ + --split_file path/to/split.npz \ + --n_samples 200 +""" +import argparse +import os +import sys +import time +import types +import importlib.util + +import numpy as np +import paddle +from scipy.spatial.distance import pdist, squareform + + +def load_model(checkpoint_path, mean, std): + """Load SchNet model from checkpoint.""" + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + for pkg in [ + "ppmat", "ppmat.models", "ppmat.models.schnet", + "ppmat.utils", "ppmat.utils.scatter", "ppmat.utils.crystal", + "ppmat.utils.paddle_aux", "ppmat.datasets", + ]: + if pkg not in sys.modules: + sys.modules[pkg] = types.ModuleType(pkg) + + for mod_name, rel_path in [ + ("ppmat.utils.paddle_aux", "ppmat/utils/paddle_aux.py"), + ("ppmat.utils.scatter", "ppmat/utils/scatter.py"), + ("ppmat.utils.crystal", "ppmat/utils/crystal.py"), + ("ppmat.models.schnet.schnet", "ppmat/models/schnet/schnet.py"), + ]: + spec = importlib.util.spec_from_file_location( + mod_name, os.path.join(repo_root, rel_path) + ) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + + from ppmat.models.schnet.schnet import SchNet + state_dict = paddle.load(checkpoint_path) + model = SchNet( + n_atom_basis=128, n_interactions=6, n_gaussians=50, cutoff=10.0, + data_mean=mean, data_std=std, readout="sum", property_names="energy_U0", + ) + model.set_state_dict(state_dict) + model.eval() + return model, state_dict + + +def evaluate(model, state_dict, qm9_dir, split_file, n_samples, cutoff=10.0): + """Evaluate model MAE on QM9 test split.""" + from ppmat.models.schnet.schnet import shifted_softplus + from ppmat.utils.scatter import scatter + + split = np.load(split_file) + test_idx = split["test_idx"] + + # Load QM9 data + try: + from ase.db import connect + db_path = os.path.join(qm9_dir, "qm9.db") if not qm9_dir.endswith(".db") else qm9_dir + db = connect(db_path) + except ImportError: + print("ase required for QM9 evaluation. Install: pip install ase") + return + + atomref = state_dict["atomref"].numpy() if "atomref" in state_dict else None + n_eval = min(n_samples, len(test_idx)) + preds, labels = [], [] + t0 = time.time() + + for count, idx in enumerate(test_idx[:n_eval]): + row = db.get(int(idx) + 1) + atoms_np = row.numbers + positions_np = row.positions + label = row.data.get("energy_U0", row.get("U0", None)) + if label is None: + continue + n_atoms = len(atoms_np) + + dists_mat = squareform(pdist(positions_np)) + src_l, dst_l, dist_l = [], [], [] + for i in range(n_atoms): + for j in range(n_atoms): + if i != j and dists_mat[i, j] < cutoff: + src_l.append(i) + dst_l.append(j) + dist_l.append(dists_mat[i, j]) + + atom_types_t = paddle.to_tensor(atoms_np, dtype="int64") + edge_src_t = paddle.to_tensor(np.array(src_l), dtype="int64") + edge_dst_t = paddle.to_tensor(np.array(dst_l), dtype="int64") + dist_t = paddle.to_tensor(np.array(dist_l, dtype=np.float32)) + batch = paddle.zeros([n_atoms], dtype="int64") + + with paddle.no_grad(): + rbf = model.rbf(dist_t) + x = model.embedding(atom_types_t) + for interaction in model.interactions: + x = interaction(x, rbf, edge_src_t, edge_dst_t, n_atoms, dist=dist_t) + x = shifted_softplus(model.output_network[0](x)) + atom_energy = model.output_network[1](x) + atom_energy = atom_energy * model.data_std + model.data_mean + if atomref is not None: + atomref_t = paddle.to_tensor(atomref) + atom_energy = atom_energy + atomref_t[atom_types_t] + energy = scatter(atom_energy, batch, dim=0, reduce="sum") + + preds.append(energy.numpy().flatten()[0]) + labels.append(float(label)) + + if (count + 1) % 100 == 0: + elapsed = time.time() - t0 + mae_so_far = np.mean(np.abs(np.array(preds) - np.array(labels))) + print(f" {count+1}/{n_eval} done ({elapsed:.1f}s), running MAE: {mae_so_far*1000:.3f} meV") + + preds = np.array(preds) + labels = np.array(labels) + mae = np.mean(np.abs(preds - labels)) + + print(f"\nResults ({len(preds)} samples):") + print(f" MAE: {mae*1000:.3f} meV/molecule") + print(f" Time: {time.time()-t0:.1f}s") + print(f" Paper reference: ~14 meV for QM9 U0") + return mae + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate SchNet on QM9 U0") + parser.add_argument("--checkpoint", required=True, help="Path to .pdparams") + parser.add_argument("--qm9_dir", required=True, help="Path to QM9 data dir") + parser.add_argument("--split_file", required=True, help="Path to split.npz") + parser.add_argument("--n_samples", type=int, default=200, help="Number of test samples") + parser.add_argument("--mean", type=float, default=-76.1160, help="Data mean") + parser.add_argument("--std", type=float, default=10.3238, help="Data std") + args = parser.parse_args() + + model, state_dict = load_model(args.checkpoint, args.mean, args.std) + evaluate(model, state_dict, args.qm9_dir, args.split_file, args.n_samples) + + +if __name__ == "__main__": + main() diff --git a/tools/test_schnet_alignment.py b/tools/test_schnet_alignment.py new file mode 100644 index 00000000..0a5c8e3e --- /dev/null +++ b/tools/test_schnet_alignment.py @@ -0,0 +1,398 @@ +# 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. + +"""Forward alignment test between SchNetPack (PyTorch) and our Paddle implementation. + +This script: +1. Loads a pretrained SchNetPack v0.3 model (PyTorch) +2. Creates a Paddle SchNet model with matching architecture +3. Transfers weights from PyTorch → Paddle +4. Constructs identical inputs for both +5. Compares forward outputs (energy predictions) + +Usage: + python tools/test_schnet_alignment.py \ + --torch_model pretrained_weights/trained_schnet_models/qm9_energy_U0/best_model \ + --paddle_params checkpoints/schnet_qm9_U0/schnet_qm9_U0.pdparams + +Requirements: + - torch, schnetpack==0.3, paddle, pgl, numpy, ase +""" + +import argparse +import importlib.util +import json +import os +import sys +import types + +import numpy as np +import paddle +import torch + +# Add project root to path +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _ROOT) + +# Pre-register ppmat packages as stubs to avoid the full import chain +# (which requires pymatgen, pgl, etc.) +for _pkg in [ + "ppmat", + "ppmat.datasets", + "ppmat.losses", + "ppmat.metrics", + "ppmat.models", + "ppmat.optimizer", + "ppmat.schedulers", + "ppmat.trainer", + "ppmat.utils", + "ppmat.sampler", +]: + if _pkg not in sys.modules: + _m = types.ModuleType(_pkg) + _m.__path__ = [] + sys.modules[_pkg] = _m + + +def _load_mod(name, relpath): + spec = importlib.util.spec_from_file_location(name, os.path.join(_ROOT, relpath)) + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +_load_mod("ppmat.utils.paddle_aux", "ppmat/utils/paddle_aux.py") +_load_mod("ppmat.utils.crystal", "ppmat/utils/crystal.py") +_load_mod("ppmat.utils.scatter", "ppmat/utils/scatter.py") +_schnet_mod = _load_mod("ppmat.models.schnet.schnet", "ppmat/models/schnet/schnet.py") + +SchNet = _schnet_mod.SchNet +GaussianRBF = _schnet_mod.GaussianRBF +shifted_softplus = _schnet_mod.shifted_softplus +from ppmat.utils.scatter import scatter + + +def create_test_molecule(): + """Create a simple water-like molecule for testing. + + Returns dict of numpy arrays representing the molecule. + """ + # 3 atoms: O(8), H(1), H(1) in a non-periodic box + atom_types = np.array([8, 1, 1], dtype=np.int64) + positions = np.array( + [ + [0.0000, 0.0000, 0.1173], + [0.0000, 0.7572, -0.4692], + [0.0000, -0.7572, -0.4692], + ], + dtype=np.float32, + ) + # Use a large box (non-periodic effectively) + lattice = np.eye(3, dtype=np.float32) * 50.0 + return { + "atom_types": atom_types, + "positions": positions, + "lattice": lattice, + } + + +def build_neighbor_list(positions, lattice, cutoff=10.0): + """Build neighbor list using distance matrix (no PBC needed for molecules). + + Returns edge_index [num_edges, 2] and pbc_offsets [num_edges, 3]. + """ + n = len(positions) + edge_src = [] + edge_dst = [] + for i in range(n): + for j in range(n): + if i != j: + dist = np.linalg.norm(positions[i] - positions[j]) + if dist < cutoff: + edge_src.append(j) + edge_dst.append(i) + edge_index = np.stack([np.array(edge_src), np.array(edge_dst)], axis=1).astype( + np.int64 + ) + pbc_offsets = np.zeros((len(edge_src), 3), dtype=np.float32) + return edge_index, pbc_offsets + + +def run_pytorch_forward(model_path, mol): + """Run SchNetPack v0.3 forward pass and return energy.""" + + # Load model (weights_only=False needed for SchNetPack v0.3 pickle format) + model = torch.load(model_path, map_location="cpu", weights_only=False) + if hasattr(model, "module"): + model = model.module + model.eval() + + # Build SchNetPack input manually (AtomsConverter broken with numpy 1.24+) + positions = mol["positions"].astype(np.float32) + atom_types = mol["atom_types"].astype(np.int64) + n_atoms = len(atom_types) + max_nbh = n_atoms - 1 # each atom sees all others + + # Build per-atom neighbor list [n_atoms, max_nbh] + nbh_list = np.full((n_atoms, max_nbh), -1, dtype=np.int64) + nbh_mask = np.zeros((n_atoms, max_nbh), dtype=np.float32) + for i in range(n_atoms): + k = 0 + for j in range(n_atoms): + if j != i and np.linalg.norm(positions[i] - positions[j]) < 10.0: + nbh_list[i, k] = j + nbh_mask[i, k] = 1.0 + k += 1 + + # cell_offset shape: [n_atoms, max_nbh, 3] (no PBC → all zeros) + cell_offset = np.zeros((n_atoms, max_nbh, 3), dtype=np.float32) + + inputs = { + "_atomic_numbers": torch.LongTensor(atom_types).unsqueeze(0), + "_positions": torch.FloatTensor(positions).unsqueeze(0), + "_cell": torch.FloatTensor(mol["lattice"]).unsqueeze(0), + "_cell_offset": torch.FloatTensor(cell_offset).unsqueeze(0), + "_neighbors": torch.LongTensor(nbh_list).unsqueeze(0), + "_neighbor_mask": torch.FloatTensor(nbh_mask).unsqueeze(0), + "_atom_mask": torch.ones(1, n_atoms), + } + + with torch.no_grad(): + result = model(inputs) + + for key in result: + if isinstance(result[key], torch.Tensor): + energy = result[key].detach().numpy() + print(f" PyTorch output['{key}'] = {energy}") + + return result + + +def run_paddle_forward( + params_path, mol, n_atom_basis=128, n_interactions=6, n_gaussians=50, cutoff=10.0 +): + """Run our Paddle SchNet forward pass and return energy.""" + # Load params to get data_mean and data_std + state_dict = paddle.load(params_path) + data_mean = float(state_dict.get("data_mean", 0.0)) + data_std = float(state_dict.get("data_std", 1.0)) + + # Create model + model = SchNet( + n_atom_basis=n_atom_basis, + n_interactions=n_interactions, + n_gaussians=n_gaussians, + cutoff=cutoff, + data_mean=data_mean, + data_std=data_std, + readout="sum", + property_names="energy_U0", + ) + + # Load weights + missing, unexpected = model.set_state_dict(state_dict) + if missing: + print(f" Missing keys: {missing}") + if unexpected: + print(f" Unexpected keys: {unexpected}") + + model.eval() + + # Build graph manually (no PGL dependency needed) + positions = mol["positions"] + atom_types = mol["atom_types"] + n_atoms = len(atom_types) + + edge_index, pbc_offsets = build_neighbor_list(positions, mol["lattice"], cutoff) + edge_src = edge_index[:, 0] + edge_dst = edge_index[:, 1] + + # Compute distances directly + dist = np.linalg.norm( + positions[edge_dst] - positions[edge_src], axis=-1 + ).astype(np.float32) + + # Convert to Paddle tensors + atom_types_t = paddle.to_tensor(atom_types, dtype="int64") + edge_src_t = paddle.to_tensor(edge_src, dtype="int64") + edge_dst_t = paddle.to_tensor(edge_dst, dtype="int64") + dist_t = paddle.to_tensor(dist) + batch = paddle.zeros([n_atoms], dtype="int64") + + with paddle.no_grad(): + rbf = model.rbf(dist_t) + x = model.embedding(atom_types_t) + for interaction in model.interactions: + x = interaction(x, rbf, edge_src_t, edge_dst_t, n_atoms, dist=dist_t) + x = shifted_softplus(model.output_network[0](x)) + atom_energy = model.output_network[1](x) # [n_atoms, 1] + + # SchNetPack standardizes per atom, then adds atomref, then sums + atom_energy = atom_energy * model.data_std + model.data_mean + + # Apply atomref if available + if "atomref" in state_dict: + atomref = paddle.to_tensor(state_dict["atomref"]) # [100, 1] + atom_energy = atom_energy + atomref[atom_types_t] + + energy = scatter(atom_energy, batch, dim=0, reduce="sum") + + energy_val = energy.numpy().flatten()[0] + print(f" Paddle energy = {energy_val:.8f}") + return {"energy": energy, "energy_val": energy_val} + + +def test_component_alignment( + model_path, params_path, n_atom_basis, n_gaussians, cutoff +): + """Test individual components match between PyTorch and Paddle.""" + print("\n=== Component-Level Alignment ===") + + # Load PyTorch model + pt_model = torch.load(model_path, map_location="cpu", weights_only=False) + if hasattr(pt_model, "module"): + pt_model = pt_model.module + pt_state = pt_model.state_dict() + + # Load Paddle params + pd_state = paddle.load(params_path) + + # Test 1: Embedding + print("\n1. Embedding weight alignment:") + pt_emb = pt_state["representation.embedding.weight"].numpy() + pd_emb = pd_state["embedding.weight"].numpy() + max_diff = np.max(np.abs(pt_emb - pd_emb)) + print(f" Max diff: {max_diff:.2e} (should be < 1e-6)") + + # Test 2: RBF offsets + print("\n2. RBF offsets alignment:") + pd_rbf = GaussianRBF(n_gaussians=n_gaussians, cutoff=cutoff) + pd_rbf_offsets = pd_rbf.offsets.numpy() + pt_rbf_offsets = np.linspace(0, cutoff, n_gaussians).astype(np.float32) + max_diff = np.max(np.abs(pd_rbf_offsets - pt_rbf_offsets)) + print(f" Max diff: {max_diff:.2e} (should be < 1e-6)") + + # Test 3: shifted_softplus + print("\n3. shifted_softplus alignment:") + test_input = paddle.to_tensor(np.array([-1.0, 0.0, 1.0, 5.0], dtype=np.float32)) + pd_ssp = shifted_softplus(test_input).numpy() + import torch.nn.functional as F + + pt_ssp = (F.softplus(torch.tensor([-1.0, 0.0, 1.0, 5.0])) - np.log(2.0)).numpy() + max_diff = np.max(np.abs(pd_ssp - pt_ssp)) + print(f" Max diff: {max_diff:.2e} (should be < 1e-6)") + + # Test 4: Filter network weight alignment (interaction 0) + print("\n4. Filter network weights (interaction 0):") + pt_fn0_w = ( + pt_state["representation.interactions.0.filter_network.0.weight"].numpy().T + ) + pd_fn0_w = pd_state["interactions.0.cfconv.filter_net.0.weight"].numpy() + max_diff = np.max(np.abs(pt_fn0_w - pd_fn0_w)) + print(f" Max diff: {max_diff:.2e} (should be < 1e-6)") + + # Test 5: Output network weights + print("\n5. Output network weights:") + pt_out0_w = pt_state["output_modules.0.out_net.1.out_net.0.weight"].numpy().T + pd_out0_w = pd_state["output_network.0.weight"].numpy() + max_diff = np.max(np.abs(pt_out0_w - pd_out0_w)) + print(f" Max diff: {max_diff:.2e} (should be < 1e-6)") + + print("\n=== Component tests done ===") + + +def main(): + parser = argparse.ArgumentParser(description="SchNet forward alignment test") + parser.add_argument( + "--torch_model", + type=str, + default="pretrained_weights/trained_schnet_models/qm9_energy_U0/best_model", + help="Path to SchNetPack v0.3 pretrained model", + ) + parser.add_argument( + "--paddle_params", + type=str, + default="checkpoints/schnet_qm9_U0/schnet_qm9_U0.pdparams", + help="Path to converted Paddle parameters", + ) + parser.add_argument( + "--n_atom_basis", type=int, default=128, help="Atom feature dimension" + ) + parser.add_argument( + "--n_interactions", type=int, default=6, help="Number of interaction blocks" + ) + parser.add_argument( + "--n_gaussians", type=int, default=50, help="Number of Gaussian RBF basis" + ) + parser.add_argument("--cutoff", type=float, default=10.0, help="Cutoff distance") + args = parser.parse_args() + + print("=" * 60) + print("SchNet Forward Alignment Test") + print("=" * 60) + + mol = create_test_molecule() + print(f"\nTest molecule: {len(mol['atom_types'])} atoms (water)") + print(f" Atoms: {mol['atom_types']}") + print(f" Positions:\n{mol['positions']}") + + # Component alignment + test_component_alignment( + args.torch_model, + args.paddle_params, + args.n_atom_basis, + args.n_gaussians, + args.cutoff, + ) + + # PyTorch forward + print("\n=== PyTorch (SchNetPack v0.3) Forward ===") + pt_result = run_pytorch_forward(args.torch_model, mol) + + # Paddle forward + print("\n=== Paddle (PaddleMaterials) Forward ===") + pd_result = run_paddle_forward( + args.paddle_params, + mol, + n_atom_basis=args.n_atom_basis, + n_interactions=args.n_interactions, + n_gaussians=args.n_gaussians, + cutoff=args.cutoff, + ) + + # Compare + print("\n" + "=" * 60) + print("COMPARISON") + print("=" * 60) + pd_energy = pd_result["energy"].numpy().flatten() + print(f" Paddle energy: {pd_energy}") + + # Extract PyTorch energy + for key, val in pt_result.items(): + if isinstance(val, torch.Tensor): + pt_energy = val.detach().numpy().flatten() + diff = np.abs(pt_energy - pd_energy) + print(f" PyTorch energy ({key}): {pt_energy}") + print(f" Abs diff: {diff}") + if np.max(diff) < 1e-4: + print(" ✅ PASS (max diff < 1e-4)") + elif np.max(diff) < 1e-2: + print(" ⚠️ WARN (max diff < 1e-2, acceptable for float32)") + else: + print(" ❌ FAIL (max diff >= 1e-2)") + + +if __name__ == "__main__": + main()