diff --git a/electronic_structure/README.md b/electronic_structure/README.md index edeece03..351e2ee7 100644 --- a/electronic_structure/README.md +++ b/electronic_structure/README.md @@ -6,28 +6,28 @@ Machine Learning Electronic Structure (MLES) is an emerging paradigm in computat ## 2.Models Matrix -| **Supported Functions** | **[InfGCN](./configs/infgcn/README.md)** | -| -------------------------------------------- | :--------: | -| **Forward Prediction · Materials Properties**| | -| Electron density | ✅ | -| **ML Capabilities · Training** | | -| Single-GPU | ✅ | -| Distributed training | ✅ | -| Mixed precision (AMP) | — | -| Fine-tuning | ✅ | -| Uncertainty / Active Learning | — | -| Dynamic→Static graphs | — | -| Compiler (CINN) opt. | — | -| **ML Capabilities · Predict** | | -| Distillation / Pruning | — | -| Standard inference | ✅ | -| Distributed inference | — | -| Compiler-level inference | — | -| **Datasets** | | -| **Materials Project** | | -| MP_EC | ✅ | -| MD17_EC | ✅ | -| QM9_EC | ✅ | -| OMol25_EC | ✅ | +| **Supported Functions** | **[InfGCN](./configs/infgcn/README.md)** | **[GPWNO](./configs/gpwno/README.md)** | +| -------------------------------------------- | :--------: | :--------: | +| **Forward Prediction · Materials Properties**| | | +| Electron density | ✅ | ✅ | +| **ML Capabilities · Training** | | | +| Single-GPU | ✅ | ✅ | +| Distributed training | ✅ | ✅ | +| Mixed precision (AMP) | — | — | +| Fine-tuning | ✅ | ✅ | +| Uncertainty / Active Learning | — | — | +| Dynamic→Static graphs | — | — | +| Compiler (CINN) opt. | — | — | +| **ML Capabilities · Predict** | | | +| Distillation / Pruning | — | — | +| Standard inference | ✅ | ✅ | +| Distributed inference | — | — | +| Compiler-level inference | — | — | +| **Datasets** | | | +| **Materials Project** | | | +| MP_EC | ✅ | ✅ | +| MD17_EC | ✅ | ✅ | +| QM9_EC | ✅ | ✅ | +| OMol25_EC | ✅ | — | -**Notice**:🌟 represent originate research work published from paddlematerials toolkit \ No newline at end of file +**Notice**:🌟 represent originate research work published from paddlematerials toolkit diff --git a/electronic_structure/configs/gpwno/README.md b/electronic_structure/configs/gpwno/README.md new file mode 100644 index 00000000..16c67bde --- /dev/null +++ b/electronic_structure/configs/gpwno/README.md @@ -0,0 +1,231 @@ +# GPWNO + +[Gaussian Plane-Wave Neural Operator for Electron Density Estimation](https://arxiv.org/abs/2402.04278) + +Official implementation: [seongsukim-ml/GPWNO](https://github.com/seongsukim-ml/GPWNO) + +## Abstract + +GPWNO is an equivariant neural operator for electron-density estimation. It combines atom-centered Gaussian basis functions with a plane-wave neural operator, allowing the model to capture both local atomic environments and global Fourier-space density patterns. The model follows the same electron-density prediction setting as InfGCN and was originally built on top of a similar data pipeline. + +
+ GPWNO overview +
+ +--- + +## Model Description + +### Overview + +Given atom types and atomic coordinates, GPWNO predicts a continuous electron-density field. The model first learns local equivariant features around atoms, then refines the field with a plane-wave operator in Fourier space. + +The main components are: + +- Gaussian atom-centered density basis +- SE(3)-equivariant coefficient learning +- Plane-wave neural operator refinement +- Optional periodic-boundary support for crystalline systems + +### Method + +#### 1) Gaussian basis expansion + +The electron density is represented with atom-centered basis functions: + +$$ +\hat{\rho}(x) = +\sum_i \sum_{n,l,m} +c_{i,nlm}\,g_n(\|x-r_i\|)\,Y_{lm}(\widehat{x-r_i}) +$$ + +where $g_n$ is a radial Gaussian basis and $Y_{lm}$ is a spherical harmonic basis. The learnable part is the coefficient tensor $c_{i,nlm}$. + +#### 2) Equivariant local message passing + +GPWNO learns basis coefficients with equivariant message passing over the atomic graph. Edge features are built from interatomic distances and angular information, preserving rotation and translation behavior needed for density prediction. + +#### 3) Plane-wave neural operator + +After local coefficient learning, GPWNO applies a Fourier-space operator to improve global consistency: + +$$ +\rho_{\mathrm{out}} = \mathcal{F}^{-1} +\left( +W(k)\,\mathcal{F}(\rho_{\mathrm{local}}) +\right) +$$ + +The Fourier module is controlled mainly by `num_fourier`, `num_fourier_time`, `width`, `fourier_mode`, `padding`, and `using_ff`. + +#### 4) Training objective and metric + +The model is trained with point-wise density regression on sampled grid points. The commonly reported metric is Normalized Mean Absolute Error: + +$$ +\mathrm{NMAE} = +\frac{\sum_i |\hat{\rho}(x_i)-\rho(x_i)|} +{\sum_i |\rho(x_i)|} +$$ + +--- + +## Dataset Description + +GPWNO uses the same electron-density datasets and dataloaders as InfGCN. + +- **MD17_EC**: Small-molecule molecular dynamics electron-density data. The supported molecules include benzene, ethanol, phenol, resorcinol, ethane, and malonaldehyde. [Data](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MD17_ES/md17_es.tar.gz). +- **QM9_EC**: QM9 molecule electron-density data stored as `qm9.json`, `qm9_data_split.json`, and `*.CHGCAR.lz4` files under `./data/data_qm9`. Data link: TBD. +- **MP_EC (cubic)**: Materials Project-style crystal electron densities stored as `*.json.xz` under `./data/data_mp`. [Data](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MP_ES/mp_es.tar), [atom dictionary](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MP_ES/crystal.json), [split file](https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MP_ES/crystal_data_split.json). + +--- + +## Key Configuration + +The PaddleMaterials configs follow the original GPWNO hyper-parameter choices: + +| Dataset | Original model profile | Main differences | +| --- | --- | --- | +| MD17_EC | `GPWNO_MD.yaml` | `num_spherical=3`, `num_fourier=40`, `padding=4`, `use_max_cell=true`, `equivariant_frame=true`, `residual=false` | +| QM9_EC | `GPWNO_QM9.yaml` | `num_spherical=4`, `num_fourier=40`, `padding=4`, `use_max_cell=true`, `equivariant_frame=true`, `residual=true` | +| MP_EC | `GPWNO_pbc.yaml` | `pbc=true`, `num_fourier=20`, `padding=0`, `use_max_cell=false`, `max_cell_size=324`, `equivariant_frame=false` | + +--- + +## Results + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Model NameDatasetNormalized MAE of DensityGPUsTraining timeConfigCheckpoint | Log
gpwno_md17_benzeneMD17_EC_Benzene3.5567% (val, 10 epochs)115hour33mingpwno_md17_benzeneTBD
gpwno_md17_ethaneMD17_EC_Ethane4.8339% (test)152hour18mingpwno_md17_ethaneTBD
gpwno_qm9QM9_EC7.4001% (test, 2 epochs)336hour33mingpwno_qm9TBD
gpwno_mpMP_EC (cubic)37.8910%137hour46mingpwno_mpTBD
+ +Note: The MD17_EC_Benzene result is reported from the epoch-10 validation of the local run `gpwno_md17_benzene_t_20260527_182931_s_42`; the corresponding checkpoint is `checkpoints/best.pdparams`. The run was intentionally recorded as a 10-epoch result. + +Note: The MD17_EC_Ethane result is reported from the final test set evaluation of the local run `gpwno_md17_ethane_t_20260529_204928_s_42`. The final validation NMAE is `4.8302%`, and the final test NMAE is `4.8339%`. + +Note: The QM9_EC result is reported from the final test set evaluation of the local run `gpwno_qm9_t_20260626_164201_s_42`, which resumed from `gpwno_qm9_t_20260625_221247_s_42/checkpoints/latest`. The best validation NMAE is `7.0865%`, and the final test NMAE is `7.4001%`. The corresponding checkpoint is `checkpoints/best.pdparams`. + +Note: The MP_EC result is reported from the final test set evaluation of the local run `gpwno_mp_resume_t_20260528_204842_s_42`. The final validation NMAE is `34.8928%`, and the final test NMAE is `37.8910%`. + +--- + +## Command + +### Data preparation + +The datasets are downloaded automatically by the dataset classes when the configured root directory is missing and `auto_download` is enabled. + +```bash +# MD17_EC will be prepared under ./data/data_md +python electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml Global.do_train=False Global.do_eval=False Global.do_test=True + +# QM9_EC will be prepared under ./data/data_qm9 after the dataset package is downloaded and extracted +python electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_qm9.yaml Global.do_train=False Global.do_eval=False Global.do_test=True +``` + +### Training + +```bash +# single-gpu training +python electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml + +# multi-gpu training +python -m paddle.distributed.launch --gpus="0,1,2,3" electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_md17_benzene.yaml + +# QM9_EC multi-gpu training +python -m paddle.distributed.launch --gpus="0,1,2" electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_qm9.yaml + +# resume training from checkpoint +python -m paddle.distributed.launch --gpus="0,1,2" electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_qm9.yaml Trainer.resume_from_checkpoint='path/to/checkpoints/latest' +``` + +### Validation + +```bash +python electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml Global.do_train=False Global.do_eval=True Global.do_test=False Trainer.pretrained_model_path='path/to/model.pdparams' + +python electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_qm9.yaml Global.do_train=False Global.do_eval=True Global.do_test=False Trainer.pretrained_model_path='path/to/checkpoints/best.pdparams' +``` + +### Testing + +```bash +python electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml Global.do_train=False Global.do_eval=False Global.do_test=True Trainer.pretrained_model_path='path/to/model.pdparams' + +python electronic_structure/train.py -c electronic_structure/configs/gpwno/gpwno_qm9.yaml Global.do_train=False Global.do_eval=False Global.do_test=True Trainer.pretrained_model_path='path/to/checkpoints/best.pdparams' +``` + +### Prediction + +```bash +python electronic_structure/predict.py -c electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml Trainer.pretrained_model_path='path/to/model.pdparams' +``` + +--- + +## Citation + +```bibtex +@inproceedings{kim2024gaussian, + title={Gaussian Plane-Wave Neural Operator for Electron Density Estimation}, + author={Kim, Seongsu and Kim, Dong Hwan and Lee, Woo Youn and Han, Seungwu}, + booktitle={International Conference on Machine Learning}, + year={2024} +} +``` + +```bibtex +@article{fandel2023infgcn, + title={InfGCN: Equivariant Neural Operator Learning with Graphon Convolution}, + author={Fandel, Dorian and von Rudorff, Guido Falk and von Lilienfeld, O. Anatole}, + journal={arXiv preprint arXiv:2311.10908}, + year={2023} +} +``` diff --git a/electronic_structure/configs/gpwno/gpwno_md17_benzene.yaml b/electronic_structure/configs/gpwno/gpwno_md17_benzene.yaml new file mode 100644 index 00000000..9108a3be --- /dev/null +++ b/electronic_structure/configs/gpwno/gpwno_md17_benzene.yaml @@ -0,0 +1,148 @@ +Global: + do_train: True + do_eval: True + do_test: True + use_voxel: False + +Trainer: + max_epochs: 32 + max_iter: 10000 + seed: 42 + output_dir: ./output/gpwno_md17_benzene + save_freq: 2000 + log_freq: 1 + start_eval_epoch: 0 + 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: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "mae" + greater_is_better: False + compute_metric_during_train: True + metric_strategy_during_eval: 'epoch' + use_visualdl: False + use_wandb: False + use_tensorboard: True + +Model: + __class_name__: GPWNO + __init_params__: + n_atom_type: 3 + num_radial: 16 + num_spherical: 3 + radial_embed_size: 64 + radial_hidden_size: 128 + num_radial_layer: 2 + num_gcn_layer: 2 + cutoff: 3.0 + grid_cutoff: 0.75 + probe_cutoff: 1.5 + is_fc: false + gauss_start: 0.5 + gauss_end: 5.0 + residual: false + pbc: false + activation: norm + num_fourier: 40 + num_fourier_time: 12 + width: 16 + fourier_mode: 0 + padding: 4 + use_max_cell: true + max_cell_size: 20 + equivariant_frame: true + normalize: true + using_ff: true + model_sharing: false + num_infgcn_layer: 3 + use_detach: false + num_spherical_RNO: 4 + scalar_mask: true + mask_cutoff: 3.0 + +Optimizer: + __class_name__: Adam + __init_params__: + lr: + __class_name__: ReduceOnPlateau + __init_params__: + learning_rate: 0.005 + factor: 0.5 + patience: 5 + min_lr: 0.00001 + by_epoch: True + indicator: "eval_loss" + indicator_name: "loss" + weight_decay: 0.0 + beta1: 0.9 + beta2: 0.999 + +Dataset: + train: + dataset: + __class_name__: SmallDensityDataset + __init_params__: + root: ./data/data_md + mol_name: benzene + split: "train" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 16 + loader: + num_workers: 4 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + val: + dataset: + __class_name__: SmallDensityDataset + __init_params__: + root: ./data/data_md + mol_name: benzene + split: "validation" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 16 + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + test: + dataset: + __class_name__: SmallDensityDataset + __init_params__: + root: ./data/data_md + mol_name: benzene + split: "test" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 1 + loader: + num_workers: 2 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + +Predict: + eval_with_no_grad: True + num_infer: null + num_vis: 2 + inf_samples: 4096 diff --git a/electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml b/electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml new file mode 100644 index 00000000..35e1d64d --- /dev/null +++ b/electronic_structure/configs/gpwno/gpwno_md17_ethane.yaml @@ -0,0 +1,148 @@ +Global: + do_train: True + do_eval: True + do_test: True + use_voxel: False + +Trainer: + max_epochs: 32 + max_iter: 10000 + seed: 42 + output_dir: ./output/gpwno_md17_ethane + save_freq: 2000 + log_freq: 1 + start_eval_epoch: 0 + 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: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "mae" + greater_is_better: False + compute_metric_during_train: True + metric_strategy_during_eval: 'epoch' + use_visualdl: False + use_wandb: False + use_tensorboard: True + +Model: + __class_name__: GPWNO + __init_params__: + n_atom_type: 3 + num_radial: 16 + num_spherical: 3 + radial_embed_size: 64 + radial_hidden_size: 128 + num_radial_layer: 2 + num_gcn_layer: 2 + cutoff: 3.0 + grid_cutoff: 0.75 + probe_cutoff: 1.5 + is_fc: false + gauss_start: 0.5 + gauss_end: 5.0 + residual: false + pbc: false + activation: norm + num_fourier: 40 + num_fourier_time: 12 + width: 16 + fourier_mode: 0 + padding: 4 + use_max_cell: true + max_cell_size: 20 + equivariant_frame: true + normalize: true + using_ff: true + model_sharing: false + num_infgcn_layer: 3 + use_detach: false + num_spherical_RNO: 4 + scalar_mask: true + mask_cutoff: 3.0 + +Optimizer: + __class_name__: Adam + __init_params__: + lr: + __class_name__: ReduceOnPlateau + __init_params__: + learning_rate: 0.005 + factor: 0.5 + patience: 5 + min_lr: 0.00001 + by_epoch: True + indicator: "eval_loss" + indicator_name: "loss" + weight_decay: 0.0 + beta1: 0.9 + beta2: 0.999 + +Dataset: + train: + dataset: + __class_name__: SmallDensityDataset + __init_params__: + root: ./data/data_md + mol_name: ethane + split: "train" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 16 + loader: + num_workers: 4 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 1024 + val: + dataset: + __class_name__: SmallDensityDataset + __init_params__: + root: ./data/data_md + mol_name: ethane + split: "validation" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 16 + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + test: + dataset: + __class_name__: SmallDensityDataset + __init_params__: + root: ./data/data_md + mol_name: ethane + split: "test" + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 1 + loader: + num_workers: 2 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + +Predict: + eval_with_no_grad: True + num_infer: null + num_vis: 2 + inf_samples: 4096 diff --git a/electronic_structure/configs/gpwno/gpwno_mp.yaml b/electronic_structure/configs/gpwno/gpwno_mp.yaml new file mode 100644 index 00000000..33e1aebe --- /dev/null +++ b/electronic_structure/configs/gpwno/gpwno_mp.yaml @@ -0,0 +1,167 @@ +Global: + do_train: True + do_eval: True + do_test: True + use_voxel: False + +Trainer: + max_epochs: 2 + max_iter: 100000 + seed: 42 + output_dir: ./output/gpwno_mp + save_freq: 2000 + log_freq: 1 + start_eval_epoch: 0 + 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: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "mae" + greater_is_better: False + compute_metric_during_train: True + metric_strategy_during_eval: 'epoch' + use_visualdl: False + use_wandb: False + use_tensorboard: True + +Model: + __class_name__: GPWNO + __init_params__: + n_atom_type: 84 + num_radial: 16 + num_spherical: 3 + radial_embed_size: 64 + radial_hidden_size: 128 + num_radial_layer: 2 + num_gcn_layer: 2 + cutoff: 3.0 + grid_cutoff: 0.75 + probe_cutoff: 1.5 + is_fc: false + gauss_start: 0.5 + gauss_end: 5.0 + residual: true + pbc: true + activation: norm + num_fourier: 20 + num_fourier_time: 12 + width: 16 + fourier_mode: 0 + padding: 0 + use_max_cell: false + max_cell_size: 324 + equivariant_frame: false + normalize: true + using_ff: true + model_sharing: false + num_infgcn_layer: 3 + use_detach: false + num_spherical_RNO: 3 + scalar_mask: true + mask_cutoff: 1.0 + +Optimizer: + __class_name__: Adam + __init_params__: + lr: + __class_name__: ReduceOnPlateau + __init_params__: + learning_rate: 0.001 + factor: 0.5 + patience: 10 + min_lr: 0.00001 + by_epoch: True + indicator: "eval_loss" + indicator_name: "loss" + weight_decay: 0.0 + beta1: 0.9 + beta2: 0.999 + +Dataset: + train: + dataset: + __class_name__: DensityDataset + __init_params__: + root: ./data/data_mp + split_file: crystal_data_split.json + atom_file: ./data/data_mp/crystal.json + extension: json + compression: xz + pbc: true + split: "train" + enable_cache: false + overwrite: false + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: True + batch_size: 16 + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 1024 + val: + dataset: + __class_name__: DensityDataset + __init_params__: + root: ./data/data_mp + split_file: crystal_data_split.json + atom_file: ./data/data_mp/crystal.json + extension: json + compression: xz + pbc: true + split: "validation" + enable_cache: false + overwrite: false + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 6 + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 4096 + test: + dataset: + __class_name__: DensityDataset + __init_params__: + root: ./data/data_mp + split_file: crystal_data_split.json + atom_file: ./data/data_mp/crystal.json + extension: json + compression: xz + pbc: true + split: "test" + rotate: false + enable_cache: false + overwrite: false + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 1 + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + +Predict: + eval_with_no_grad: True + num_infer: 100 + num_vis: 2 + inf_samples: 4096 diff --git a/electronic_structure/configs/gpwno/gpwno_qm9.yaml b/electronic_structure/configs/gpwno/gpwno_qm9.yaml new file mode 100644 index 00000000..29943cad --- /dev/null +++ b/electronic_structure/configs/gpwno/gpwno_qm9.yaml @@ -0,0 +1,167 @@ +Global: + do_train: True + do_eval: True + do_test: True + use_voxel: False + +Trainer: + max_epochs: 2 + max_iter: null + seed: 42 + output_dir: ./output/gpwno_qm9 + save_freq: 1 + log_freq: 1 + start_eval_epoch: 0 + 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: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "mae" + 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__: GPWNO + __init_params__: + n_atom_type: 5 + num_radial: 16 + num_spherical: 4 + radial_embed_size: 64 + radial_hidden_size: 128 + num_radial_layer: 2 + num_gcn_layer: 2 + cutoff: 3.0 + grid_cutoff: 0.75 + probe_cutoff: 1.5 + is_fc: false + gauss_start: 0.5 + gauss_end: 5.0 + residual: true + pbc: false + activation: norm + num_fourier: 40 + num_fourier_time: 12 + width: 16 + fourier_mode: 0 + padding: 4 + use_max_cell: true + max_cell_size: 20 + equivariant_frame: true + normalize: true + using_ff: true + model_sharing: false + num_infgcn_layer: 3 + use_detach: false + num_spherical_RNO: 4 + scalar_mask: true + mask_cutoff: 3.0 + +Optimizer: + __class_name__: Adam + __init_params__: + lr: + __class_name__: ReduceOnPlateau + __init_params__: + learning_rate: 0.005 + factor: 0.5 + patience: 5 + min_lr: 0.00001 + by_epoch: True + indicator: "eval_loss" + indicator_name: "loss" + weight_decay: 0.000000000001 + beta1: 0.9 + beta2: 0.999 + +Dataset: + train: + dataset: + __class_name__: DensityDataset + __init_params__: + root: ./data/data_qm9 + split_file: qm9_data_split.json + atom_file: ./data/data_qm9/qm9.json + extension: CHGCAR + compression: lz4 + pbc: false + split: "train" + enable_cache: false + overwrite: false + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: True + batch_size: 32 + loader: + num_workers: 4 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 1024 + val: + dataset: + __class_name__: DensityDataset + __init_params__: + root: ./data/data_qm9 + split_file: qm9_data_split.json + atom_file: ./data/data_qm9/qm9.json + extension: CHGCAR + compression: lz4 + pbc: false + split: "validation" + enable_cache: false + overwrite: false + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + loader: + num_workers: 4 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + test: + dataset: + __class_name__: DensityDataset + __init_params__: + root: ./data/data_qm9 + split_file: qm9_data_split.json + atom_file: ./data/data_qm9/qm9.json + extension: CHGCAR + compression: lz4 + pbc: false + split: "test" + rotate: false + enable_cache: false + overwrite: false + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 8 + loader: + num_workers: 4 + use_shared_memory: False + collate_fn: DensityCollator + collate_params: + n_samples: 2048 + +Predict: + eval_with_no_grad: True + num_infer: 100 + num_vis: 2 + inf_samples: 4096 diff --git a/electronic_structure/docs/gpwno_overview.png b/electronic_structure/docs/gpwno_overview.png new file mode 100644 index 00000000..973e6e9a Binary files /dev/null and b/electronic_structure/docs/gpwno_overview.png differ diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py index 95d73232..b25a5fc9 100644 --- a/ppmat/models/__init__.py +++ b/ppmat/models/__init__.py @@ -41,6 +41,7 @@ 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.gpwno.GPWNO import GPWNO from ppmat.models.mateno.mateno import MatENO from ppmat.utils import download from ppmat.utils import logger @@ -66,6 +67,7 @@ "DiffPrior", "DiffNMR", "InfGCN", + "GPWNO", "MatENO", ] diff --git a/ppmat/models/gpwno/GPWNO.py b/ppmat/models/gpwno/GPWNO.py new file mode 100644 index 00000000..a03cd1d1 --- /dev/null +++ b/ppmat/models/gpwno/GPWNO.py @@ -0,0 +1,912 @@ +# 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. + +import sys +import numpy as np +import paddle +from ..common.e3nn import o3 +from ..common.e3nn.math import soft_one_hot_linspace +from ..common.e3nn.nn import Activation, Extract, FullyConnectedNet +from .GPWNO_utils import * +from ..common.orbital import GaussianOrbital +from .PWNO_utils import * +from ppmat.datasets.graph_utils.infgcn_graph_utils import radius +from ppmat.datasets.graph_utils.infgcn_graph_utils import radius_graph +# from ....paddle_geometric.paddle_geometric.nn import radius, radius_graph +from paddle_scatter.scatter import scatter + + +def pbc_vec(vec, cell): + """ + Apply periodic boundary condition to the vector + :param vec: original vector of (N, K, 3) + :param cell: cell frame of (N, 3, 3) + :return: shortest vector of (N, K, 3) + """ + coord = vec @ paddle.linalg.inv(x=cell) + coord = coord - paddle.round(coord) + pbc_vec = coord @ cell + return pbc_vec.detach(), coord.detach() + + +class PWNO(paddle.nn.Module): + def __init__( + self, + modes1, + modes2, + modes3, + width, + num_fourier_time, + padding=0, + num_layers=2, + using_ff=False, + ): + super(PWNO, self).__init__() + """ + The overall network. It contains 4 layers of the Fourier layer. + 1. Lift the input to the desire channel dimension by self.fc0 . + 2. 4 layers of the integral operators u' = (W + K)(u). + W defined by self.w; K defined by self.conv . + 3. Project from the channel space to the output space by self.fc1 and self.fc2 . + + input: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y), x, y, t). It's a constant function in time, except for the last index. + input shape: (batchsize, x=64, y=64, t=40, c=13) + output: the solution of the next 40 timesteps + output shape: (batchsize, x=64, y=64, t=40, c=1) + """ + self.modes1 = modes1 + self.modes2 = modes2 + self.modes3 = modes3 + self.width = width + self.num_fourier_time = num_fourier_time + self.padding = padding + self.fc0 = paddle.compat.nn.Linear(self.num_fourier_time + 3, self.width) + self.num_layers = num_layers + if using_ff: + self.conv = paddle.nn.ModuleList( + [ + SpectralConv3d_FFNO( + self.width, self.width, self.modes1, self.modes2, self.modes3 + ) + for _ in range(self.num_layers) + ] + ) + else: + self.conv = paddle.nn.ModuleList( + [ + SpectralConv3d( + self.width, self.width, self.modes1, self.modes2, self.modes3 + ) + for _ in range(self.num_layers) + ] + ) + self.w = paddle.nn.ModuleList( + [ + paddle.nn.Conv3d(self.width, self.width, 1) + for _ in range(self.num_layers) + ] + ) + self.bn = paddle.nn.ModuleList( + [ + paddle.nn.BatchNorm3D(num_features=self.width) + for _ in range(self.num_layers) + ] + ) + + def forward(self, residue, fourier_grid): + x = paddle.cat([residue, fourier_grid], dim=-1) + x = self.fc0(x) + x = x.permute(0, 4, 1, 2, 3) + if self.padding != 0: + x = paddle.compat.nn.functional.pad( + x, [0, self.padding, 0, self.padding, 0, self.padding] + ) + for i in range(self.num_layers): + x1 = self.conv[i](x) + x2 = self.w[i](x) + x = x1 + x2 + if i != self.num_layers - 1: + x = paddle.nn.functional.gelu(x) + if self.padding != 0: + x = x[..., :, : -self.padding, : -self.padding, : -self.padding] + x = x.permute(0, 2, 3, 4, 1) + return x + + +class GPWNO(paddle.nn.Layer): + def __init__( + self, + n_atom_type, + num_radial, + num_spherical, + radial_embed_size, + radial_hidden_size, + num_fourier=64, + num_fourier_time=4, + num_radial_layer=2, + num_gcn_layer=3, + cutoff=3.0, + grid_cutoff=3.0, + probe_cutoff=3.0, + is_fc=True, + gauss_start=0.5, + gauss_end=5.0, + activation="norm", + residual=True, + pbc=False, + width=20, + padding=6, + use_max_cell=True, + max_cell_size=22, + equivariant_frame=True, + probe_and_node=False, + product=False, + normalize=True, + model_sharing=True, + num_infgcn_layer=3, + input_infgcn=False, + using_ff=True, + scalar_mask=True, + use_detach=True, + mask_cutoff=3.0, + scalar_inv=False, + num_spherical_RNO=None, + positive_output=False, + atomic_gauss_dist=False, + input_dist=False, + atom_info=None, + fourier_mode=0, + *args, + **kwargs, + ): + """ + Implement the GPWNO model for electron density estimation + :param n_atom_type: number of atom types + :param num_radial: number of radial basis + :param num_spherical: maximum number of spherical harmonics for each radial basis, + number of spherical basis will be (num_spherical + 1)^2 + :param radial_embed_size: embedding size of the edge length + :param radial_hidden_size: hidden size of the radial network + :param num_radial_layer: number of hidden layers in the radial network + :param num_gcn_layer: number of GPWNO layers + :param cutoff: cutoff distance for building the molecular graph + :param grid_cutoff: cutoff distance for building the grid-atom graph + :param is_fc: whether the GPWNO layer should use fully connected tensor product + :param gauss_start: start coefficient of the Gaussian radial basis + :param gauss_end: end coefficient of the Gaussian radial basis + :param activation: activation type for the GPWNO layer, can be ['scalar', 'norm'] + :param residual: whether to use the residue prediction layer + :param pbc: whether the data satisfy the periodic boundary condition + """ + super(GPWNO, self).__init__(*args, **kwargs) + self.n_atom_type = n_atom_type + self.num_radial = num_radial + self.num_spherical = num_spherical + self.radial_embed_size = radial_embed_size + self.radial_hidden_size = radial_hidden_size + self.num_radial_layer = num_radial_layer + self.num_gcn_layer = num_gcn_layer + self.cutoff = cutoff + self.grid_cutoff = grid_cutoff + self.is_fc = is_fc + self.gauss_start = gauss_start + self.gauss_end = gauss_end + self.activation = activation + self.residual = residual + self.pbc = pbc + self.num_fourier = num_fourier + self.num_fourier_time = num_fourier_time + self.probe_cutoff = probe_cutoff + self.use_max_cell = use_max_cell + self.max_cell_size = max_cell_size + self.equivariant_frame = equivariant_frame + self.probe_and_node = probe_and_node + self.normalize = normalize + self.model_sharing = model_sharing + self.num_infgcn_layer = num_infgcn_layer + self.input_infgcn = input_infgcn + self.scalar_mask = scalar_mask + self.use_detach = use_detach + self.mask_cutoff = mask_cutoff + self.scalar_inv = scalar_inv + self.positive_output = positive_output + self.atomic_gauss_dist = atomic_gauss_dist + self.input_dist = input_dist + self.atom_info = atom_info + self.fourier_mode = fourier_mode + self.num_spherical_RNO = num_spherical_RNO + if model_sharing == True or num_spherical_RNO is None: + self.num_spherical_RNO = num_spherical + self.product = product + assert activation in ["scalar", "norm"] + self.embedding = paddle.nn.Embedding(n_atom_type, num_radial) + self.irreps_sh = o3.Irreps.spherical_harmonics(num_spherical, p=1) + self.irreps_feat = (self.irreps_sh * num_radial).sort().irreps.simplify() + self.irreps_sh_RNO = o3.Irreps.spherical_harmonics(num_spherical_RNO, p=1) + self.irreps_feat_RNO = ( + (self.irreps_sh_RNO * num_radial).sort().irreps.simplify() + ) + self.gcns_RNO = paddle.nn.ModuleList( + [ + GCNLayer( + f"{num_radial}x0e" if i == 0 else self.irreps_feat_RNO, + self.irreps_feat_RNO, + self.irreps_sh_RNO, + radial_embed_size, + num_radial_layer, + radial_hidden_size, + is_fc=is_fc, + **kwargs, + ) + for i in range(num_gcn_layer) + ] + ) + if self.activation == "scalar": + self.act = ScalarActivation( + self.irreps_feat, paddle.nn.functional.silu, paddle.sigmoid + ) + self.act_RNO = ScalarActivation( + self.irreps_feat_RNO, paddle.nn.functional.silu, paddle.sigmoid + ) + else: + self.act = NormActivation(self.irreps_feat) + self.act_RNO = NormActivation(self.irreps_feat_RNO) + if self.model_sharing == False: + self.infgcns = paddle.nn.ModuleList( + [ + GCNLayer( + f"{num_radial}x0e" if i == 0 else self.irreps_feat, + self.irreps_feat, + self.irreps_sh, + radial_embed_size, + num_radial_layer, + radial_hidden_size, + is_fc=is_fc, + **kwargs, + ) + for i in range(num_infgcn_layer) + ] + ) + self.residue = None + if self.residual: + self.residue = GCNLayer( + self.irreps_feat, + "0e", + self.irreps_sh, + radial_embed_size, + num_radial_layer, + radial_hidden_size, + is_fc=True, + use_sc=False, + **kwargs, + ) + self.probe_gcn = GCNLayer( + self.irreps_feat_RNO, + f"{num_fourier_time}x0e", + self.irreps_sh_RNO, + radial_embed_size, + num_radial_layer, + radial_hidden_size, + is_fc=True, + use_sc=False, + **kwargs, + ) + self.act_probe = NormActivation(f"{num_fourier_time}x0e") + self.mode = ( + self.num_fourier // 2 if self.fourier_mode == 0 else self.fourier_mode + ) + self.width = width + if self.input_infgcn: + self.num_fourier_time += 1 + if self.input_dist: + self.num_fourier_time += 1 + if self.atomic_gauss_dist: + self.num_fourier_time += 10 + self.PWNO = PWNO( + modes1=self.mode, + modes2=self.mode, + modes3=self.mode, + width=self.width, + num_fourier_time=self.num_fourier_time, + padding=padding, + using_ff=using_ff, + ) + self.scalar_field_gcn = GCNLayer( + f"{self.width}x0e", + f"0e", + self.irreps_sh_RNO, + radial_embed_size, + num_radial_layer, + radial_hidden_size, + is_fc=True, + use_sc=False, + **kwargs, + ) + self.orbital = GaussianOrbital( + gauss_start, gauss_end, num_radial, num_spherical + ) + if self.atomic_gauss_dist and atom_info is not None: + import json + + with open(atom_info) as f: + atom_info = json.load(f) + atom_radius = [info["radius"] for info in atom_info] + self.atom_radius = [radius for idx, radius in enumerate(atom_radius)] + self.atom_radius = paddle.FloatTensor(self.atom_radius) + + def forward(self, batch_idx): + """ + ppmat-style forward + + batch_idx: + graph: PGL graph, with x / pos / batch_idx + grid_coord: [B, K, 3] + density: optional, [B, K] + density_mask: optional, [B, K] + infos: optional, list[dict], may contain 'cell' + """ + graph = batch_idx["graph"] + grid = batch_idx["grid_coord"] + infos = batch_idx.get("infos", None) + + density_gt = batch_idx.get("density", None) + mask = batch_idx.get("density_mask", None) + + pred, aux = self._forward_density( + atom_types=graph.x, + atom_coord=graph.pos, + grid=grid, + batch_idx=graph.batch, + infos=infos, + ) + + loss_dict = {} + if density_gt is not None: + loss, mae = self._compute_loss(pred, density_gt, mask) + loss_dict["loss"] = loss + loss_dict["mae"] = mae + + return { + "loss_dict": loss_dict, + "pred_dict": {"density": pred}, + "aux_dict": aux, + } + + def _compute_loss(self, pred, label, mask=None): + if mask is not None: + pred = pred * mask + label = label * mask + denom = paddle.sum(mask) + 1e-8 + loss = paddle.sum((pred - label) ** 2) / denom + mae = paddle.sum(paddle.abs(pred - label)) / ( + paddle.sum(label) + 1e-8 + ) + else: + loss = paddle.mean((pred - label) ** 2) + mae = paddle.sum(paddle.abs(pred - label)) / ( + paddle.sum(label) + 1e-8 + ) + return loss, mae + + def _forward_density(self,atom_types,atom_coord,grid,batch_idx,infos,): + + device = atom_coord.place + + cell = None + if infos is not None and "cell" in infos[0]: + cell = paddle.stack([info["cell"] for info in infos], axis=0).to(device) + + atom_types = atom_types.to(device) + grid = grid.to(device) + + feat = self.embedding(atom_types) + n_graph, n_sample = grid.shape[0], grid.shape[1] + + if self.use_max_cell and self.equivariant_frame: + atom_center = scatter(atom_coord, batch_idx, dim=0, reduce="mean") + atom_coord = atom_coord - atom_center[batch_idx] + grid -= atom_center.unsqueeze(1) + edge_index = radius_graph(atom_coord, self.cutoff, batch_idx, loop=False) + src, dst = edge_index + edge_vec = atom_coord[src] - atom_coord[dst] + edge_len = edge_vec.norm(dim=-1) + 1e-08 + edge_feat = o3.spherical_harmonics( + list(range(self.num_spherical + 1)), + edge_vec / edge_len[..., None], + normalize=False, + normalization="integral", + ) + edge_embed = soft_one_hot_linspace( + edge_len, + start=0.0, + end=self.cutoff, + number=self.radial_embed_size, + basis="gaussian", + cutoff=False, + ).mul(self.radial_embed_size**0.5) + edge_feat_RNO = o3.spherical_harmonics( + list(range(self.num_spherical_RNO + 1)), + edge_vec / edge_len[..., None], + normalize=False, + normalization="integral", + ) + if self.model_sharing == False: + infgcn_feat = feat + for i, gcn in enumerate(self.infgcns): + infgcn_feat = gcn( + edge_index, + infgcn_feat, + edge_feat, + edge_embed, + dim_size=atom_types.size(0), + ) + if i != self.num_infgcn_layer - 1: + infgcn_feat = self.act(infgcn_feat) + for i, gcn in enumerate(self.gcns_RNO): + feat = gcn( + edge_index, feat, edge_feat_RNO, edge_embed, dim_size=atom_types.size(0) + ) + if i != self.num_gcn_layer - 1: + feat = self.act_RNO(feat) + bins_lin = paddle.linspace(0, 1, self.num_fourier).to(atom_coord.place) + if self.pbc: + half = len(bins_lin) // 2 + super_bins = paddle.cat( + [bins_lin[half:-1] - 1, bins_lin, bins_lin[1:half] + 1] + ).to(atom_coord.place) + super_bins = paddle.meshgrid(super_bins, super_bins, super_bins) + super_probe = paddle.stack(super_bins, dim=-1).to(atom_coord.place) + bins_idx = paddle.arange(len(bins_lin)) + bins_idx[-1] = 0 + super_bins_idx = paddle.cat( + [bins_idx[half:-1], bins_idx, bins_idx[1:half]] + ).to(atom_coord.place) + smart_idx = ( + paddle.arange((len(bins_idx) - 1) ** 3) + .reshape(len(bins_idx) - 1, len(bins_idx) - 1, len(bins_idx) - 1) + .to(atom_coord.place) + ) + super_probe_idx_help = paddle.stack( + paddle.meshgrid(super_bins_idx, super_bins_idx, super_bins_idx), dim=-1 + ).reshape(-1, 3) + super_probe_idx = smart_idx[ + super_probe_idx_help[:, 0], + super_probe_idx_help[:, 1], + super_probe_idx_help[:, 2], + ].reshape(-1) + del super_probe_idx_help + if self.use_max_cell and self.equivariant_frame: + bins_lin -= 0.5 + bins = paddle.meshgrid(bins_lin, bins_lin, bins_lin) + probe = paddle.stack(bins, dim=-1).to(atom_coord.place) + cell_inp = cell + if self.use_max_cell: + if self.max_cell_size is None: + self.max_cell_size = 20 + if self.equivariant_frame: + num_batch = grid.size(0) + new_cell = [] + for i in range(num_batch): + atom_bat = atom_coord[batch_idx == i] + atoms_centered = atom_bat - atom_bat.mean(dim=0) + R = paddle.matmul(atoms_centered.t(), atoms_centered) + _, vec = paddle.linalg.eigh(x=R) + vec /= paddle.linalg.norm(vec, dim=0) + new_cell.append(vec) + new_cell = paddle.stack(new_cell, dim=0) + max_cell = new_cell * self.max_cell_size + else: + max_cell_tensor = np.eyes(3) * self.max_cell_size + max_cell = paddle.FloatTensor(self.max_cell_size).to(atom_coord.place) + max_cell = max_cell.unsqueeze(0).repeat(grid.size(0), 1, 1) + cell_inp = max_cell + probe = paddle.einsum("ijkl,blm->bijkm", probe, cell_inp).detach() + probe_log = probe.cpu() + if self.use_max_cell and self.equivariant_frame: + probe_log += atom_center.reshape(-1, 1, 1, 1, 3).cpu() + probe_flat = probe.reshape(-1, 3) + probe_batch = paddle.arange(grid.size(0), device=grid.device).repeat_interleave( + len(bins_lin) ** 3 + ) + if self.pbc: + super_probe = paddle.einsum( + "ijkl,blm->bijkm", super_probe, cell_inp + ).detach() + super_probe_flat = super_probe.reshape(-1, 3) + super_probe_batch = paddle.arange( + grid.size(0), device=grid.device + ).repeat_interleave(len(super_bins_idx) ** 3) + super_probe_idx_batch = paddle.arange( + grid.size(0), device=grid.device + ).repeat_interleave(len(super_probe_idx)) * len( + bins_lin + ) ** 3 + super_probe_idx.repeat( + grid.size(0) + ) + if self.pbc: + super_probe_dst, probe_src = radius( + atom_coord, + super_probe_flat, + self.probe_cutoff, + batch_idx, + super_probe_batch, + ) + probe_dst = super_probe_idx_batch[super_probe_dst] + else: + probe_dst, probe_src = radius( + atom_coord, probe_flat, self.probe_cutoff, batch_idx, probe_batch + ) + avg_probe_degree = ( + probe_dst.size(0) / probe_flat.size(0) if probe_flat.size(0) != 0 else 0 + ) + avg_probe_exist_degree = ( + probe_dst.size(0) / probe_dst.unique().size(0) + if probe_dst.unique().size(0) != 0 + else 0 + ) + avg_probe_exist_rate = ( + probe_dst.unique().size(0) / probe_flat.size(0) + if probe_dst.unique().size(0) != 0 + else 0 + ) + if self.pbc: + probe_edge = super_probe_flat[super_probe_dst] - atom_coord[probe_src] + else: + probe_edge = probe_flat[probe_dst] - atom_coord[probe_src] + probe_len = paddle.norm(probe_edge, dim=-1) + 1e-08 + probe_edge_feat = o3.spherical_harmonics( + list(range(self.num_spherical_RNO + 1)), + probe_edge / (probe_len[..., None] + 1e-08), + normalize=False, + normalization="integral", + ) + probe_edge_embed = soft_one_hot_linspace( + probe_len, + start=0.0, + end=self.probe_cutoff, + number=self.radial_embed_size, + basis="gaussian", + cutoff=False, + ).mul(self.radial_embed_size**0.5) + probe_feat = self.probe_gcn.forward_mean( + (probe_src, probe_dst), + feat, + probe_edge_feat, + probe_edge_embed, + dim_size=probe_flat.size(0), + ) + probe_feat = self.act_probe(probe_feat) + probe_feat = probe_feat.reshape( + grid.size(0), self.num_fourier, self.num_fourier, self.num_fourier, -1 + ) + minimal_dist_grid = paddle.zeros(grid.size(0), grid.size(1)) + for b_idx in range(grid.size(0)): + if self.pbc: + minimal_dist_grid[b_idx] = paddle.compat.min( + paddle.compat.min( + paddle.norm( + grid[b_idx].unsqueeze(1).unsqueeze(1) + - atom_coord[batch_idx == b_idx].unsqueeze(0).unsqueeze(0) + + paddle.stack( + [ + ( + i * cell[b_idx][0] + + j * cell[b_idx][1] + + k * cell[b_idx][2] + ) + for i in [-1, 0, 1] + for j in [-1, 0, 1] + for k in [-1, 0, 1] + ], + dim=0, + ).reshape(1, 27, 1, 3), + dim=-1, + ), + dim=-2, + )[0], + dim=-1, + )[0] + else: + minimal_dist_grid[b_idx] = paddle.compat.min( + paddle.norm( + grid[b_idx].unsqueeze(1) + - atom_coord[batch_idx == b_idx].unsqueeze(0), + dim=-1, + ), + dim=-1, + )[0] + mask = minimal_dist_grid > self.mask_cutoff + atomic_dist_probe = paddle.zeros(grid.size(0), self.num_fourier**3, 10).to( + atom_coord.place + ) + alpha = paddle.tensor([0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0]).to( + atom_coord.place + ) + alpha = alpha.reshape(1, 1, 10) + minimal_dist_probe = paddle.zeros(grid.size(0), self.num_fourier**3).to( + atom_coord.place + ) + for b_idx in range(grid.size(0)): + if self.pbc: + pr = probe.reshape(grid.size(0), -1, 3) + dist_min = paddle.compat.min( + paddle.norm( + pr[b_idx].unsqueeze(1).unsqueeze(1) + - atom_coord[batch_idx == b_idx].unsqueeze(0).unsqueeze(0) + + paddle.stack( + [ + ( + i * cell[b_idx][0] + + j * cell[b_idx][1] + + k * cell[b_idx][2] + ) + for i in [-1, 0, 1] + for j in [-1, 0, 1] + for k in [-1, 0, 1] + ], + dim=0, + ).reshape(1, 27, 1, 3), + dim=-1, + ), + dim=-2, + )[0] + if self.atomic_gauss_dist and self.atom_info is not None: + self.atom_radius = self.atom_radius.to(atom_coord.place) + atomic_dist_probe[b_idx] = paddle.sum( + paddle.exp( + -alpha + * ( + dist_min + / self.atom_radius[ + atom_types[batch_idx == b_idx] + ].unsqueeze(0) + ).unsqueeze(-1) + ), + dim=-2, + ) + minimal_dist_probe[b_idx] = paddle.compat.min(dist_min, dim=-1)[0] + if self.input_infgcn: + probe_vec = probe.reshape(grid.size(0), -1, 3)[ + batch_idx + ] - atom_coord.unsqueeze(-2) + if self.pbc: + probe_vec, frac = pbc_vec(probe_vec, cell[batch_idx]) + density_res = paddle.zeros(grid.size(0), self.num_fourier**3, 1).to( + atom_coord.place + ) + for b_idx in range(grid.size(0)): + if self.model_sharing == False: + density_res[b_idx] = ( + ( + self.orbital(probe_vec[b_idx]) + * infgcn_feat[batch_idx == b_idx].unsqueeze(0) + ) + .sum(dim=-1) + .reshape(-1, 1) + ) + else: + density_res[b_idx] = ( + ( + self.orbital(probe_vec[b_idx]) + * feat[batch_idx == b_idx].unsqueeze(0) + ) + .sum(dim=-1) + .reshape(-1, 1) + ) + density_res = density_res.reshape( + grid.size(0), self.num_fourier, self.num_fourier, self.num_fourier, -1 + ) + if self.model_sharing == True: + infgcn_feat = feat + if self.residual and self.input_infgcn: + res_feat = infgcn_feat + probe_edge_feat_res = o3.spherical_harmonics( + list(range(self.num_spherical + 1)), + probe_edge / (probe_len[..., None] + 1e-08), + normalize=False, + normalization="integral", + ) + residue = self.residue( + (probe_src, probe_dst), + res_feat, + probe_edge_feat_res, + probe_edge_embed, + dim_size=probe_flat.size(0), + ) + density_res = density_res + residue.reshape( + grid.size(0), + self.num_fourier, + self.num_fourier, + self.num_fourier, + -1, + ) + probe_feat = paddle.cat([probe_feat, density_res], dim=-1) + if self.input_dist: + minimal_dist_probe = paddle.exp(-minimal_dist_probe * 0.5).unsqueeze(-1) + minimal_dist_probe = minimal_dist_probe.reshape( + grid.size(0), self.num_fourier, self.num_fourier, self.num_fourier, 1 + ) + probe_feat = paddle.cat([probe_feat, minimal_dist_probe], dim=-1) + if self.atomic_gauss_dist: + atomic_dist_probe = atomic_dist_probe.reshape( + grid.size(0), self.num_fourier, self.num_fourier, self.num_fourier, 10 + ) + probe_feat = paddle.cat([probe_feat, atomic_dist_probe], dim=-1) + fourier_grid = self.get_grid(probe_feat.shape, probe_feat.device) + probe_feat = self.PWNO(probe_feat, fourier_grid) + probe_feat = probe_feat.reshape(-1, probe_feat.size(-1)) + sample_flat = grid.reshape(-1, 3) + sample_batch = ( + paddle.arange(n_graph, device=grid.device) + .repeat_interleave(n_sample) + .detach() + ) + probe_and_node = probe_flat + probe_and_node_batch = probe_batch + if self.pbc: + sample_dst, super_probe_and_node_src = radius( + super_probe_flat, + sample_flat, + self.grid_cutoff, + super_probe_batch, + sample_batch, + ) + probe_and_node_src = super_probe_idx_batch[super_probe_and_node_src] + else: + sample_dst, probe_and_node_src = radius( + probe_and_node, + sample_flat, + self.grid_cutoff, + probe_and_node_batch, + sample_batch, + ) + avg_sample_degree = ( + sample_dst.size(0) / sample_flat.size(0) if sample_flat.size(0) != 0 else 0 + ) + avg_sample_exist_degree = ( + sample_dst.size(0) / sample_dst.unique().size(0) + if sample_dst.unique().size(0) != 0 + else 0 + ) + avg_sample_exist_rate = ( + sample_dst.unique().size(0) / sample_flat.size(0) + if sample_dst.unique().size(0) != 0 + else 0 + ) + if self.pbc: + sample_edge = ( + sample_flat[sample_dst] - super_probe_flat[super_probe_and_node_src] + ) + else: + sample_edge = sample_flat[sample_dst] - probe_flat[probe_and_node_src] + sample_len = paddle.norm(sample_edge, dim=-1) + 1e-08 + sample_edge_feat = o3.spherical_harmonics( + list(range(self.num_spherical_RNO + 1)), + sample_edge / (sample_len[..., None] + 1e-08), + normalize=False, + normalization="integral", + ) + sample_edge_embed = soft_one_hot_linspace( + sample_len, + start=0.0, + end=self.grid_cutoff, + number=self.radial_embed_size, + basis="gaussian", + cutoff=False, + ).mul(self.radial_embed_size**0.5) + probe_and_node_feat = probe_feat + scalar_field_gcn = self.scalar_field_gcn.forward_mean( + (probe_and_node_src, sample_dst), + probe_and_node_feat, + sample_edge_feat, + sample_edge_embed, + dim_size=sample_flat.size(0), + ) + if self.residual: + grid_flat = grid.view(-1, 3) + grid_batch = paddle.arange(n_graph, device=grid.device).repeat_interleave( + n_sample + ) + grid_dst, node_src = radius( + atom_coord, grid_flat, self.cutoff, batch_idx, grid_batch + ) + grid_edge = grid_flat[grid_dst] - atom_coord[node_src] + grid_len = paddle.norm(grid_edge, dim=-1) + 1e-08 + grid_edge_feat = o3.spherical_harmonics( + list(range(self.num_spherical + 1)), + grid_edge / (grid_len[..., None] + 1e-08), + normalize=False, + normalization="integral", + ) + grid_edge_embed = soft_one_hot_linspace( + grid_len, + start=0.0, + end=self.cutoff, + number=self.radial_embed_size, + basis="gaussian", + cutoff=False, + ).mul(self.radial_embed_size**0.5) + res_feat = feat + if self.model_sharing == False: + res_feat = infgcn_feat + residue = self.residue( + (node_src, grid_dst), + res_feat, + grid_edge_feat, + grid_edge_embed, + dim_size=grid_flat.size(0), + ) + else: + residue = 0.0 + sample_vec = grid[batch_idx] - atom_coord.unsqueeze(-2) + if self.pbc: + sample_vec, frac = pbc_vec(sample_vec, cell[batch_idx]) + orbital = self.orbital(sample_vec) + if self.model_sharing == False: + density = (orbital * infgcn_feat.unsqueeze(1)).sum(dim=-1) + else: + density = (orbital * feat.unsqueeze(1)).sum(dim=-1) + density = scatter(density, batch_idx, dim=0, reduce="sum") + #scalar_field = scalar_field_gcn.reshape(grid.size(0), grid.size(1)).real() + scalar_field = scalar_field_gcn.reshape([grid.shape[0], grid.shape[1]]) + if self.residual: + density = density + residue.view(*density.size()) + if self.scalar_mask: + # Match the mask dtype and device with the scalar field. + mask_float = mask.astype(paddle.float32).to(scalar_field.place) + scalar_field = scalar_field * mask_float + + if self.scalar_inv: + mask_inv_float = (1 - mask_float).astype(paddle.float32) + density = density * mask_inv_float + if self.positive_output: + density_res = density + scalar_field + density_res[(density_res < 0) & (density > 0) & (scalar_field < 0)] = 0 + scalar_field[(density_res < 0) & (density > 0) & (scalar_field < 0)] = 0 + density = density_res + else: + density = density + scalar_field + scalar_field_influence = (scalar_field.abs().sum() / density.sum()).detach() + + aux = { + "avg_probe_degree": avg_probe_degree, + "avg_probe_exist_degree": avg_probe_exist_degree, + "avg_probe_exist_rate": avg_probe_exist_rate, + "avg_sample_degree": avg_sample_degree, + "avg_sample_exist_degree": avg_sample_exist_degree, + "avg_sample_exist_rate": avg_sample_exist_rate, + "scalar_field": scalar_field, + "coefficient_field": density - scalar_field, + "scalar_field_influence": scalar_field_influence, + "probe": probe, + } + + return density, aux + + + def get_grid(self, shape, device): + batchsize, size_x, size_y, size_z = shape[0], shape[1], shape[2], shape[3] + gridx = paddle.linspace(0, 1, size_x) + gridx = gridx.reshape(1, size_x, 1, 1, 1).repeat( + [batchsize, 1, size_y, size_z, 1] + ) + gridy = paddle.linspace(0, 1, size_y) + gridy = gridy.reshape(1, 1, size_y, 1, 1).repeat( + [batchsize, size_x, 1, size_z, 1] + ) + gridz = paddle.linspace(0, 1, size_z) + gridz = gridz.reshape(1, 1, 1, size_z, 1).repeat( + [batchsize, size_x, size_y, 1, 1] + ) + return paddle.cat((gridx, gridy, gridz), dim=-1).to(device) diff --git a/ppmat/models/gpwno/GPWNO_utils.py b/ppmat/models/gpwno/GPWNO_utils.py new file mode 100644 index 00000000..ec568107 --- /dev/null +++ b/ppmat/models/gpwno/GPWNO_utils.py @@ -0,0 +1,212 @@ +# 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. + +import paddle +from ..common.e3nn import o3 +from ..common.e3nn.nn import Activation, Extract, FullyConnectedNet +# from torch_scatter import scatter +from paddle_scatter.scatter import scatter + + + +class ScalarActivation(paddle.nn.Module): + """ + Use the invariant scalar features to gate higher order equivariant features. + Adapted from `e3nn.nn.Gate`. + """ + + def __init__(self, irreps_in, act_scalars, act_gates): + """ + :param irreps_in: input representations + :param act_scalars: scalar activation function + :param act_gates: gate activation function (for higher order features) + """ + super(ScalarActivation, self).__init__() + self.irreps_in = o3.Irreps(irreps_in) + self.num_spherical = len(self.irreps_in) + irreps_scalars = self.irreps_in[0:1] + irreps_gates = irreps_scalars * (self.num_spherical - 1) + irreps_gated = self.irreps_in[1:] + self.act_scalars = Activation(irreps_scalars, [act_scalars]) + self.act_gates = Activation( + irreps_gates, [act_gates] * (self.num_spherical - 1) + ) + self.extract = Extract( + self.irreps_in, + [irreps_scalars, irreps_gated], + instructions=[(0,), tuple(range(1, self.irreps_in.lmax + 1))], + ) + self.mul = o3.ElementwiseTensorProduct(irreps_gates, irreps_gated) + + def forward(self, features): + scalars, gated = self.extract(features) + scalars_out = self.act_scalars(scalars) + if gated.shape[-1]: + gates = self.act_gates(scalars.repeat(1, self.num_spherical - 1)) + gated_out = self.mul(gates, gated) + features = paddle.cat([scalars_out, gated_out], dim=-1) + else: + features = scalars_out + return features + + +class NormActivation(paddle.nn.Module): + """ + Use the norm of the higher order equivariant features to gate themselves. + Idea from the TFN paper. + """ + + def __init__( + self, + irreps_in, + act_scalars=paddle.nn.functional.silu, + act_vectors=paddle.sigmoid, + ): + """ + :param irreps_in: input representations + :param act_scalars: scalar activation function + :param act_vectors: vector activation function (for the norm of higher order features) + """ + super(NormActivation, self).__init__() + self.irreps_in = o3.Irreps(irreps_in) + self.scalar_irreps = self.irreps_in[0:1] + self.vector_irreps = self.irreps_in[1:] + self.act_scalars = act_scalars + self.act_vectors = act_vectors + self.scalar_idx = self.irreps_in[0].mul + inner_out = o3.Irreps([(mul, (0, 1)) for mul, _ in self.vector_irreps]) + self.inner_prod = o3.TensorProduct( + self.vector_irreps, + self.vector_irreps, + inner_out, + [(i, i, i, "uuu", False) for i in range(len(self.vector_irreps))], + ) + self.mul = o3.ElementwiseTensorProduct(inner_out, self.vector_irreps) + + def forward(self, features): + scalars = self.act_scalars(features[..., : self.scalar_idx]) + vectors = features[..., self.scalar_idx :] + norm = paddle.sqrt(self.inner_prod(vectors, vectors) + 1e-08) + act = self.act_vectors(norm) + vectors_out = self.mul(act, vectors) + return paddle.cat([scalars, vectors_out], dim=-1) + + +class GCNLayer(paddle.nn.Module): + def __init__( + self, + irreps_in, + irreps_out, + irreps_edge, + radial_embed_size, + num_radial_layer, + radial_hidden_size, + is_fc=True, + use_sc=True, + irrep_normalization="component", + path_normalization="element", + *args, + **kwargs + ): + """ + A single InfGCN layer for Tensor Product-based message passing. + If the tensor product is fully connected, we have (for every path) + + .. math:: + z_w=\\sum_{uv}w_{uvw}x_u\\otimes y_v=\\sum_{u}w_{uw}x_u \\otimes y + + Else, we have + + .. math:: + z_u=x_u\\otimes \\sum_v w_{uv}y_v=w_u (x_u\\otimes y) + + Here, uvw are radial (channel) indices of the first input, second input, and output, respectively. + Notice that in our model, the second input is always the spherical harmonics of the edge vector, + so the index v can be safely ignored. + + :param irreps_in: irreducible representations of input node features + :param irreps_out: irreducible representations of output node features + :param irreps_edge: irreducible representations of edge features + :param radial_embed_size: embedding size of the edge length + :param num_radial_layer: number of hidden layers in the radial network + :param radial_hidden_size: hidden size of the radial network + :param is_fc: whether to use fully connected tensor product + :param use_sc: whether to use self-connection + :param irrep_normalization: representation normalization passed to the `o3.FullyConnectedTensorProduct` + :param path_normalization: path normalization passed to the `o3.FullyConnectedTensorProduct` + """ + super(GCNLayer, self).__init__() + self.irreps_in = o3.Irreps(irreps_in) + self.irreps_out = o3.Irreps(irreps_out) + self.irreps_edge = o3.Irreps(irreps_edge) + self.radial_embed_size = radial_embed_size + self.num_radial_layer = num_radial_layer + self.radial_hidden_size = radial_hidden_size + self.is_fc = is_fc + self.use_sc = use_sc + if self.is_fc: + self.tp = o3.FullyConnectedTensorProduct( + self.irreps_in, + self.irreps_edge, + self.irreps_out, + internal_weights=False, + shared_weights=False, + irrep_normalization=irrep_normalization, + path_normalization=path_normalization, + ) + else: + instr = [ + (i_1, i_2, i_out, "uvu", True) + for i_1, (_, ir_1) in enumerate(self.irreps_in) + for i_2, (_, ir_edge) in enumerate(self.irreps_edge) + for i_out, (_, ir_out) in enumerate(self.irreps_out) + if ir_out in ir_1 * ir_edge + ] + self.tp = o3.TensorProduct( + self.irreps_in, + self.irreps_edge, + self.irreps_out, + instr, + internal_weights=False, + shared_weights=False, + irrep_normalization=irrep_normalization, + path_normalization=path_normalization, + ) + self.fc = FullyConnectedNet( + [radial_embed_size] + + num_radial_layer * [radial_hidden_size] + + [self.tp.weight_numel], + paddle.nn.functional.silu, + ) + self.sc = None + if self.use_sc: + self.sc = o3.Linear(self.irreps_in, self.irreps_out) + + def forward(self, edge_index, node_feat, edge_feat, edge_embed, dim_size=None): + src, dst = edge_index + weight = self.fc(edge_embed) + out = self.tp(node_feat[src], edge_feat, weight=weight) + out = scatter(out, dst, dim=0, dim_size=dim_size, reduce="sum") + if self.use_sc: + out = out + self.sc(node_feat) + return out + + def forward_mean(self, edge_index, node_feat, edge_feat, edge_embed, dim_size=None): + src, dst = edge_index + weight = self.fc(edge_embed) + out = self.tp(node_feat[src], edge_feat, weight=weight) + out = scatter(out, dst, dim=0, dim_size=dim_size, reduce="mean") + if self.use_sc: + out = out + self.sc(node_feat) + return out diff --git a/ppmat/models/gpwno/PWNO_utils.py b/ppmat/models/gpwno/PWNO_utils.py new file mode 100644 index 00000000..0ab5d59b --- /dev/null +++ b/ppmat/models/gpwno/PWNO_utils.py @@ -0,0 +1,175 @@ +# 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. + +import numpy as np +import paddle + + +class SpectralConv3d(paddle.nn.Module): + def __init__(self, in_channels, out_channels, modes1, modes2, modes3): + super(SpectralConv3d, self).__init__() + """ + 3D Fourier layer. It does FFT, linear transform, and Inverse FFT. + """ + self.in_channels = in_channels + self.out_channels = out_channels + self.modes1 = modes1 + self.modes2 = modes2 + self.modes3 = modes3 + self.scale = 1 / (in_channels * out_channels) + self.weights1 = paddle.nn.Parameter( + self.scale + * paddle.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + ) + self.weights2 = paddle.nn.Parameter( + self.scale + * paddle.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + ) + self.weights3 = paddle.nn.Parameter( + self.scale + * paddle.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + ) + self.weights4 = paddle.nn.Parameter( + self.scale + * paddle.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=paddle.complex64, + ) + ) + + def compl_mul3d(self, input, weights): + return paddle.einsum("bixyz,ioxyz->boxyz", input, weights) + + def forward(self, x): + batchsize = x.shape[0] + # Pass explicit dimensions for the 3D real FFT. + x_ft = paddle.fft.rfftn(x, dim=[-3, -2, -1], norm="ortho") + out_ft = paddle.zeros( + [ + batchsize, + self.out_channels, + x.shape[-3], + x.shape[-2], + x.shape[-1] // 2 + 1, + ], + dtype=paddle.complex64, + device=x.device, + ) + out_ft[:, :, : self.modes1, : self.modes2, : self.modes3] = self.compl_mul3d( + x_ft[:, :, : self.modes1, : self.modes2, : self.modes3], self.weights1 + ) + out_ft[:, :, -self.modes1 :, : self.modes2, : self.modes3] = self.compl_mul3d( + x_ft[:, :, -self.modes1 :, : self.modes2, : self.modes3], self.weights2 + ) + out_ft[:, :, : self.modes1, -self.modes2 :, : self.modes3] = self.compl_mul3d( + x_ft[:, :, : self.modes1, -self.modes2 :, : self.modes3], self.weights3 + ) + out_ft[:, :, -self.modes1 :, -self.modes2 :, : self.modes3] = self.compl_mul3d( + x_ft[:, :, -self.modes1 :, -self.modes2 :, : self.modes3], self.weights4 + ) + # Use the original spatial shape for the inverse 3D real FFT. + x = paddle.fft.irfftn( + out_ft, + s=(x.shape[-3], x.shape[-2], x.shape[-1]), + dim=[-3, -2, -1], + norm="ortho", + ) + return x + + +class SpectralConv3d_FFNO(paddle.nn.Module): + def __init__(self, in_channels, out_channels, modes1, modes2, modes3): + super(SpectralConv3d_FFNO, self).__init__() + """ + 3D Fourier layer for FFNO-style factorized spectral convolution. + """ + self.in_channels = in_channels + self.out_channels = out_channels + self.modes_x = modes1 + self.modes_y = modes2 + self.modes_z = modes3 + self.fourier_weight = paddle.nn.ParameterList(parameters=[]) + for n_modes in [self.modes_x, self.modes_y, self.modes_z]: + weight = paddle.randn([in_channels, out_channels, n_modes, 2], dtype=paddle.float32) + paddle.nn.init.xavier_normal_(weight) + self.fourier_weight.append(paddle.nn.Parameter(weight)) + + def forward(self, x): + B, I, S1, S2, S3 = x.shape # [batch, in_ch, x, y, z] + + # Spectral convolution along the z axis. + x_ftz = paddle.fft.rfftn(x, dim=[-1,], norm="ortho") + out_ft = x_ftz.new_zeros(B, I, S1, S2, S3 // 2 + 1) + out_ft[:, :, :, :, : self.modes_z] = paddle.einsum( + "bixyz,ioz->boxyz", + x_ftz[:, :, :, :, : self.modes_z], + paddle.view_as_complex(self.fourier_weight[2]) + ) + xz = paddle.fft.irfft(out_ft, n=S3, dim=-1, norm="ortho") + + # Spectral convolution along the y axis. + x_fty = paddle.fft.rfftn(x, dim=[-2,], norm="ortho") + out_ft = x_ftz.new_zeros(B, I, S1, S2 // 2 + 1, S3) + out_ft[:, :, :, : self.modes_y, :] = paddle.einsum( + "bixyz,ioy->boxyz", + x_fty[:, :, :, : self.modes_y, :], + paddle.view_as_complex(self.fourier_weight[1]) + ) + xy = paddle.fft.irfft(out_ft, n=S2, dim=-2, norm="ortho") + + # Spectral convolution along the x axis. + x_ftx = paddle.fft.rfftn(x, dim=[-3,], norm="ortho") + out_ft = x_ftz.new_zeros(B, I, S1 // 2 + 1, S2, S3) + out_ft[:, :, : self.modes_x, :, :] = paddle.einsum( + "bixyz,iox->boxyz", + x_ftx[:, :, : self.modes_x, :, :], + paddle.view_as_complex(self.fourier_weight[0]) + ) + xx = paddle.fft.irfft(out_ft, n=S1, dim=-3, norm="ortho") + + x = xx + xy + xz + return x + + +def plane_wave_sum(x): + N = x.shape[-1] + n = paddle.arange(N).float().to(x.device) + k = n.unsqueeze(1) + M = paddle.exp(-2.0j * np.pi * k * n) + return paddle.matmul(M, x)