+
+
+ | Model Name |
+ Dataset |
+ Property |
+ MAE |
+ GPUs |
+ Training time |
+ Config |
+ Checkpoint | Log |
+
+
+
+
+ | spherenet_md17_aspirin |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.26 / 0.44 |
+ 1 |
+ ~6 h |
+ config |
+ checkpoint | log |
+
+
+ | spherenet_md17_benzene_old |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.14 / 0.21 |
+ 1 |
+ ~3 h |
+ config |
+ checkpoint | log |
+
+
+ | spherenet_md17_ethanol |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.10 / 0.23 |
+ 1 |
+ ~2 h |
+ config |
+ checkpoint | log |
+
+
+ | spherenet_md17_malonaldehyde |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.17 / 0.32 |
+ 1 |
+ ~2 h |
+ config |
+ checkpoint | log |
+
+
+ | spherenet_md17_naphthalene |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.16 / 0.26 |
+ 1 |
+ ~5 h |
+ config |
+ checkpoint | log |
+
+
+ | spherenet_md17_salicylic |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.22 / 0.38 |
+ 1 |
+ ~5 h |
+ config |
+ checkpoint | log |
+
+
+ | spherenet_md17_toluene |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.12 / 0.21 |
+ 1 |
+ ~4 h |
+ config |
+ checkpoint | log |
+
+
+ | spherenet_md17_uracil |
+ MD17 |
+ Energy (kcal/mol) / Force (kcal/mol/Å) |
+ 0.12 / 0.30 |
+ 1 |
+ ~2 h |
+ config |
+ checkpoint | log |
+
+
+
+
+### Training
+
+```bash
+# Single-GPU training — MD17 aspirin (energy + force)
+python interatomic_potentials/train.py \
+ -c interatomic_potentials/configs/spherenet/spherenet_md17_aspirin.yaml
+```
+
+### Validation
+
+```bash
+python interatomic_potentials/train.py \
+ -c interatomic_potentials/configs/spherenet/spherenet_md17_aspirin.yaml \
+ Global.do_eval=True Global.do_train=False Global.do_test=False \
+ Trainer.pretrained_model_path='your_model.pdparams'
+```
+
+### Testing
+
+```bash
+python interatomic_potentials/train.py \
+ -c interatomic_potentials/configs/spherenet/spherenet_md17_aspirin.yaml \
+ Global.do_test=True Global.do_train=False Global.do_eval=False \
+ Trainer.pretrained_model_path='your_model.pdparams'
+```
+
+### Prediction
+
+```bash
+# Molecular prediction
+python interatomic_potentials/predict.py \
+ --model_name spherenet_md17_aspirin \
+ --xyz_file_path ./interatomic_potentials/example_data/xyz/md17_aspirin.xyz
+
+# Using a local checkpoint
+python interatomic_potentials/predict.py \
+ --config_path ./interatomic_potentials/configs/spherenet/spherenet_md17_aspirin.yaml \
+ --checkpoint_path ./output/spherenet_aspirin/checkpoints/best.pdparams \
+ --xyz_file_path ./interatomic_potentials/example_data/xyz/md17_aspirin.xyz
+```
+
+## Citation
+
+```bibtex
+@inproceedings{liu2021spherenet,
+ title={Spherical Message Passing for 3D Molecular Graphs},
+ author={Liu, Yi and Wang, Limei and Liu, Meng and Lin, Yuchao and Zhang, Xuan and
+ Oztekin, Bora and Ji, Shuiwang},
+ booktitle={International Conference on Learning Representations (ICLR)},
+ year={2021}
+}
+```
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_aspirin.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_aspirin.yaml
new file mode 100644
index 00000000..114b7cd0
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_aspirin.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_aspirin
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: aspirin
+ split: train
+ split_file: ./data/md17/splits/aspirin_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: aspirin
+ split: val
+ split_file: ./data/md17/splits/aspirin_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: aspirin
+ split: test
+ split_file: ./data/md17/splits/aspirin_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_benzene_old.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_benzene_old.yaml
new file mode 100644
index 00000000..d3e976d5
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_benzene_old.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_benzene_old
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: benzene_old
+ split: train
+ split_file: ./data/md17/splits/benzene_old_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: benzene_old
+ split: val
+ split_file: ./data/md17/splits/benzene_old_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: benzene_old
+ split: test
+ split_file: ./data/md17/splits/benzene_old_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_ethanol.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_ethanol.yaml
new file mode 100644
index 00000000..6d1cc97c
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_ethanol.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_ethanol
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: ethanol
+ split: train
+ split_file: ./data/md17/splits/ethanol_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: ethanol
+ split: val
+ split_file: ./data/md17/splits/ethanol_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: ethanol
+ split: test
+ split_file: ./data/md17/splits/ethanol_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_malonaldehyde.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_malonaldehyde.yaml
new file mode 100644
index 00000000..f0a2353d
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_malonaldehyde.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_malonaldehyde
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: malonaldehyde
+ split: train
+ split_file: ./data/md17/splits/malonaldehyde_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: malonaldehyde
+ split: val
+ split_file: ./data/md17/splits/malonaldehyde_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: malonaldehyde
+ split: test
+ split_file: ./data/md17/splits/malonaldehyde_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_naphthalene.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_naphthalene.yaml
new file mode 100644
index 00000000..1443cf56
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_naphthalene.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_naphthalene
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: naphthalene
+ split: train
+ split_file: ./data/md17/splits/naphthalene_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: naphthalene
+ split: val
+ split_file: ./data/md17/splits/naphthalene_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: naphthalene
+ split: test
+ split_file: ./data/md17/splits/naphthalene_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_salicylic.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_salicylic.yaml
new file mode 100644
index 00000000..9d95dd8e
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_salicylic.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_salicylic
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: salicylic
+ split: train
+ split_file: ./data/md17/splits/salicylic_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: salicylic
+ split: val
+ split_file: ./data/md17/splits/salicylic_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: salicylic
+ split: test
+ split_file: ./data/md17/splits/salicylic_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_toluene.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_toluene.yaml
new file mode 100644
index 00000000..0b98f28b
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_toluene.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_toluene
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: toluene
+ split: train
+ split_file: ./data/md17/splits/toluene_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: toluene
+ split: val
+ split_file: ./data/md17/splits/toluene_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: toluene
+ split: test
+ split_file: ./data/md17/splits/toluene_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/configs/spherenet/spherenet_md17_uracil.yaml b/interatomic_potentials/configs/spherenet/spherenet_md17_uracil.yaml
new file mode 100644
index 00000000..06610164
--- /dev/null
+++ b/interatomic_potentials/configs/spherenet/spherenet_md17_uracil.yaml
@@ -0,0 +1,148 @@
+Global:
+ label_names:
+ - energy
+ do_train: true
+ do_eval: true
+ do_test: true
+
+ build_molecule_cfg:
+ format: dict
+ sanitize: false
+ add_hs: false
+ remove_hs: false
+ kekulize: false
+ num_cpus: 10
+
+ graph_converter:
+ __class_name__: RadiusGraphConverter
+ __init_params__:
+ cutoff: 5.0
+ return_triplet_indices: true
+ num_cpus: 10
+
+Trainer:
+ max_epochs: 500
+ seed: 42
+ output_dir: ./output/spherenet_uracil
+ save_freq: 50
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ 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
+ compute_metric_during_train: false
+ metric_strategy_during_eval: step
+ pretrained_model_path:
+Model:
+ __class_name__: SphereNet
+ __init_params__:
+ energy_and_force: true
+ cutoff: 5.0
+ num_layers: 4
+ hidden_channels: 128
+ out_channels: 1
+ int_emb_size: 64
+ basis_emb_size_dist: 8
+ basis_emb_size_angle: 8
+ basis_emb_size_torsion: 8
+ out_emb_channels: 256
+ num_spherical: 7
+ num_radial: 6
+ envelope_exponent: 5
+ num_before_skip: 1
+ num_after_skip: 2
+ num_output_layers: 3
+ output_init: zeros
+ property_name: energy
+
+Metric:
+ energy:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+ force:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+Optimizer:
+ __class_name__: Adam
+ __init_params__:
+ beta1: 0.9
+ beta2: 0.999
+ lr:
+ __class_name__: Cosine
+ __init_params__:
+ learning_rate: 0.0005
+ eta_min: 0.00001
+ by_epoch: true
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: uracil
+ split: train
+ split_file: ./data/md17/splits/uracil_train_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: true
+ drop_last: true
+ val:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: uracil
+ split: val
+ split_file: ./data/md17/splits/uracil_val_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: MD17Dataset
+ __init_params__:
+ path: ./data/md17
+ name: uracil
+ split: test
+ split_file: ./data/md17/splits/uracil_test_idx.npy
+ build_molecule_cfg: ${Global.build_molecule_cfg}
+ build_graph_cfg: ${Global.graph_converter}
+ force_key: force
+ loader:
+ num_workers: 4
+ use_shared_memory: false
+ collate_fn: RadiusGraphCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 8
+ shuffle: false
+ drop_last: false
+
+Predict:
+ eval_with_no_grad: false
+ graph_converter: ${Global.graph_converter}
diff --git a/interatomic_potentials/docs/SphereNet.png b/interatomic_potentials/docs/SphereNet.png
new file mode 100644
index 00000000..0e5d6d37
Binary files /dev/null and b/interatomic_potentials/docs/SphereNet.png differ
diff --git a/interatomic_potentials/example_data/xyz/md17_aspirin.xyz b/interatomic_potentials/example_data/xyz/md17_aspirin.xyz
new file mode 100644
index 00000000..51a736e4
--- /dev/null
+++ b/interatomic_potentials/example_data/xyz/md17_aspirin.xyz
@@ -0,0 +1,23 @@
+21
+MD17 aspirin inference example
+C -2.595942 -2.269642 0.066255
+C -1.718569 -1.136781 -0.375382
+O -1.568786 -0.790663 -1.540230
+O -1.142709 -0.550219 0.746129
+C -0.235077 0.460871 0.417163
+C -0.698267 1.777324 0.510780
+C 0.164026 2.833783 0.224236
+C 1.484579 2.575250 -0.141719
+C 1.949765 1.258119 -0.215915
+C 1.094915 0.182227 0.070449
+C 1.590339 -1.218640 0.021824
+O 0.977083 -2.225822 0.323594
+O 2.864388 -1.301050 -0.408129
+H -3.006067 -2.773570 -0.813588
+H -3.422512 -1.884425 0.668231
+H -2.008791 -2.994227 0.636116
+H -1.726984 1.977713 0.796840
+H -0.193200 3.859031 0.284092
+H 2.156660 3.399997 -0.367989
+H 2.986228 1.084400 -0.496432
+H 3.048921 -2.263677 -0.406327
diff --git a/interatomic_potentials/predict.py b/interatomic_potentials/predict.py
index cd671d17..891b575a 100644
--- a/interatomic_potentials/predict.py
+++ b/interatomic_potentials/predict.py
@@ -24,6 +24,7 @@
from pymatgen.core import Structure
from tqdm import tqdm
+from ppmat.datasets.build_molecule import BuildMolecule
from ppmat.datasets.transform import build_post_transforms
from ppmat.models import build_graph_converter
from ppmat.models import build_model
@@ -35,8 +36,8 @@
class PotentialPredictor:
"""Potential predictor.
- This class provides an interface for predicting properties of crystalline
- structures using pre-trained deep learning models. Supports two initialization
+ This class provides an interface for predicting properties of crystal structures and
+ molecules using pre-trained deep learning models. Supports two initialization
modes:
1. **Automatic Model Loading**
@@ -101,7 +102,7 @@ def __init__(
self.model.eval()
- predict_config = config.get("Predict", None)
+ predict_config = config.get("Predict") or {}
self.predict_config = predict_config
self.eval_with_no_grad = predict_config.get("eval_with_no_grad", True)
@@ -183,6 +184,39 @@ def from_cif_file(self, cif_file_path, save_path=None):
return result
+ def from_xyz_file(self, xyz_file_path, save_path=None):
+ """Predict molecular energy and forces from one XYZ file."""
+ if save_path is not None:
+ assert save_path.endswith(".csv"), "save_path must end with .csv"
+ if not osp.isfile(xyz_file_path) or not xyz_file_path.endswith(".xyz"):
+ raise ValueError(f"Expected one XYZ file, but got: {xyz_file_path}")
+ if self.graph_converter_fn is None:
+ raise ValueError("Molecular prediction requires a graph converter.")
+
+ molecule = BuildMolecule(format="xyz_file", sanitize=False)(xyz_file_path)
+ if molecule is None:
+ raise ValueError(f"Failed to parse XYZ file: {xyz_file_path}")
+ graph = self.graph_converter_fn(molecule)
+
+ if self.eval_with_no_grad:
+ with paddle.no_grad():
+ result = self.model.predict(graph)
+ else:
+ result = self.model.predict(graph)
+ result = self.post_process(result)
+
+ if save_path is not None:
+ row = {"xyz_file": osp.basename(xyz_file_path)}
+ row.update(
+ {
+ key: value.tolist() if hasattr(value, "tolist") else value
+ for key, value in result.items()
+ }
+ )
+ pd.DataFrame([row]).to_csv(save_path, index=False)
+ logger.info(f"Saved the prediction result to {save_path}")
+ return result
+
if __name__ == "__main__":
@@ -214,9 +248,15 @@ def from_cif_file(self, cif_file_path, save_path=None):
argparse.add_argument(
"--cif_file_path",
type=str,
- default="./interatomic_potentials/",
+ default=None,
help="Path to the CIF file whose material properties you want to predict.",
)
+ argparse.add_argument(
+ "--xyz_file_path",
+ type=str,
+ default=None,
+ help="Path to the XYZ molecule whose energy and forces to predict.",
+ )
argparse.add_argument(
"--save_path",
type=str,
@@ -232,5 +272,10 @@ def from_cif_file(self, cif_file_path, save_path=None):
checkpoint_path=args.checkpoint_path,
)
- results = predictor.from_cif_file(args.cif_file_path, args.save_path)
+ if args.xyz_file_path is not None:
+ results = predictor.from_xyz_file(args.xyz_file_path, args.save_path)
+ elif args.cif_file_path is not None:
+ results = predictor.from_cif_file(args.cif_file_path, args.save_path)
+ else:
+ raise ValueError("Provide --xyz_file_path or --cif_file_path.")
print(results)
diff --git a/ppmat/datasets/__init__.py b/ppmat/datasets/__init__.py
index 05d3de49..b6703993 100644
--- a/ppmat/datasets/__init__.py
+++ b/ppmat/datasets/__init__.py
@@ -31,6 +31,7 @@
from ppmat.datasets import collate_fn
from ppmat.datasets.high_level_water_dataset import HighLevelWaterDataset
+from ppmat.datasets.gmtnet_dataset import GMTNetDielectricDataset
from ppmat.datasets.jarvis_dataset import JarvisDataset
from ppmat.datasets.matbench_dataset import MatbenchDataset
from ppmat.datasets.mp20_dataset import AlexMP20MatterGenDataset
@@ -42,11 +43,12 @@
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.small_density_dataset import SmallDensityDataset
from ppmat.datasets.sfin_dataset import SFINDataset
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.md17_dataset import MD17Dataset # noqa
from ppmat.datasets.omol25_dataset import OMol25Dataset
from ppmat.datasets.split_mptrj_data import none_to_zero
from ppmat.datasets.transform import build_transforms
@@ -63,12 +65,14 @@
"MPTrjDataset",
"JarvisDataset",
"HighLevelWaterDataset",
+ "GMTNetDielectricDataset",
"MSDnmrDataset",
"MatbenchDataset",
"DensityDataset",
"SmallDensityDataset",
"SFINDataset",
"OMol25Dataset",
+ "MD17Dataset",
]
INFO_CLASS_REGISTRY: Dict[str, type] = {
diff --git a/ppmat/datasets/build_molecule.py b/ppmat/datasets/build_molecule.py
index 32755fdd..945c2d9d 100644
--- a/ppmat/datasets/build_molecule.py
+++ b/ppmat/datasets/build_molecule.py
@@ -23,14 +23,16 @@
from p_tqdm import p_map
from rdkit import Chem
+from rdkit.Geometry import Point3D
class BuildMolecule:
"""Build RDKit Mol from different formats.
Args:
- format (Literal["smiles","mol_block","mol_file","sdf_file", "inchi","dict",
- "rdmol"]): format of input molecules data used by convertion of RDKit
+ format (Literal["smiles","mol_block","mol_file","sdf_file","xyz_block",
+ "xyz_file","inchi","dict","rdmol"]): format of input molecules data used
+ by convertion of RDKit
sanitize (bool): Whether to sanitize the molecule using RDKit after construction
(e.g., validate valence, adjust bond orders). Defaults to True.
add_hs (bool): Whether to add explicit hydrogen atoms to the molecule.
@@ -47,7 +49,15 @@ class BuildMolecule:
def __init__(
self,
format: Literal[
- "smiles", "mol_block", "mol_file", "sdf_file", "inchi", "dict", "rdmol"
+ "smiles",
+ "mol_block",
+ "mol_file",
+ "sdf_file",
+ "xyz_block",
+ "xyz_file",
+ "inchi",
+ "dict",
+ "rdmol",
],
sanitize: bool = True,
add_hs: bool = False,
@@ -101,13 +111,25 @@ def build_one(
elif format == "sdf_file":
suppl = Chem.SDMolSupplier(str(mol_data), sanitize=sanitize, removeHs=False)
mol = next((m for m in suppl if m is not None), None)
+ elif format == "xyz_block":
+ mol = Chem.MolFromXYZBlock(str(mol_data))
+ elif format == "xyz_file":
+ mol = Chem.MolFromXYZFile(str(mol_data))
elif format == "inchi":
mol = Chem.MolFromInchi(str(mol_data))
elif format == "dict":
- mol_block = mol_data.get("mol_block", None)
- if mol_block is None:
- raise ValueError("dict format requires key 'mol_block'.")
- mol = Chem.MolFromMolBlock(mol_block, sanitize=sanitize)
+ atomic_numbers = mol_data["atomic_numbers"]
+ positions = mol_data["positions"]
+ mol = Chem.RWMol()
+ for atomic_number in atomic_numbers:
+ mol.AddAtom(Chem.Atom(int(atomic_number)))
+ mol = mol.GetMol()
+ conformer = Chem.Conformer(len(atomic_numbers))
+ for atom_idx, position in enumerate(positions):
+ conformer.SetAtomPosition(
+ atom_idx, Point3D(*(float(value) for value in position))
+ )
+ mol.AddConformer(conformer)
elif format == "rdmol":
mol = mol_data
else:
diff --git a/ppmat/datasets/collate_fn.py b/ppmat/datasets/collate_fn.py
index 9073af4c..88ad46fa 100644
--- a/ppmat/datasets/collate_fn.py
+++ b/ppmat/datasets/collate_fn.py
@@ -88,6 +88,68 @@ def __call__(self, batch: List[Any]) -> Any:
)
+class RadiusGraphCollator:
+ """Collates radius-graph samples for SphereNet.
+
+ Each sample is expected to contain a ``pgl.Graph`` under ``graph``. The graph
+ comes from ``RadiusGraphConverter`` and stores cached SphereNet triplet
+ indices in ``edge_feat``. Graph fields remain inside the batched graph and
+ are unpacked by the model. Variable-length sample fields should be marked by
+ the dataset with ``ConcatData`` and are collated by ``DefaultCollator``.
+ """
+
+ def __call__(self, batch):
+ graphs = [sample["graph"] for sample in batch]
+ num_edges_list = [np.asarray(graph.edges).shape[0] for graph in graphs]
+ num_triplets_list = [
+ np.asarray(graph.edge_feat["ti_idx_kj"]).shape[0]
+ for graph in graphs
+ ]
+
+ edge_offsets = np.cumsum([0] + num_edges_list[:-1])
+ triplet_offsets = np.cumsum([0] + num_triplets_list[:-1])
+
+ triplet_fields = {
+ key: []
+ for key in ("idx_kj", "idx_ji", "idx_lk", "idx_triplet")
+ }
+ for i, graph in enumerate(graphs):
+ edge_feat = graph.edge_feat
+ triplet_fields["idx_kj"].append(
+ np.asarray(edge_feat["ti_idx_kj"], dtype=np.int64)
+ + edge_offsets[i]
+ )
+ triplet_fields["idx_ji"].append(
+ np.asarray(edge_feat["ti_idx_ji"], dtype=np.int64)
+ + edge_offsets[i]
+ )
+ triplet_fields["idx_lk"].append(
+ np.asarray(edge_feat["ti_idx_lk"], dtype=np.int64)
+ + edge_offsets[i]
+ )
+ triplet_fields["idx_triplet"].append(
+ np.asarray(edge_feat["ti_idx_triplet"], dtype=np.int64)
+ + triplet_offsets[i]
+ )
+
+ graph = pgl.Graph.batch(graphs)
+ graph.edge_feat.update(
+ {
+ f"ti_{key}": np.concatenate(value)
+ for key, value in triplet_fields.items()
+ }
+ )
+
+ result = DefaultCollator()(
+ [
+ {key: value for key, value in sample.items() if key != "graph"}
+ for sample in batch
+ ]
+ )
+ result["graph"] = graph
+ return result
+
+
class DensityCollator:
def __init__(
self,
diff --git a/ppmat/datasets/gmtnet_dataset.py b/ppmat/datasets/gmtnet_dataset.py
new file mode 100644
index 00000000..ca911164
--- /dev/null
+++ b/ppmat/datasets/gmtnet_dataset.py
@@ -0,0 +1,462 @@
+# 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.
+"""GMTNet dielectric dataset backed by the normalized frozen pickle."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import operator
+import pickle
+from collections.abc import Mapping
+from contextlib import nullcontext
+from dataclasses import dataclass
+from importlib import resources
+from pathlib import Path
+from typing import Any, ClassVar
+
+import numpy as np
+import paddle
+from paddle.io import Dataset
+from pymatgen.core import Structure
+
+from ppmat.models.gmtnet.gmtnet_graph_converter import GMTNetGraphConverter
+
+_RECORD_COUNT = 4713
+_SPLIT_SIZES = {"train": 3770, "val": 471, "test": 472}
+_DEFAULT_SPLIT_RESOURCE = (
+ "configs",
+ "gmtnet",
+ "split_gmtnet_dielectric_seed32.json",
+)
+_CONVERTER_DEFAULTS = {
+ "cutoff": 4.0,
+ "max_neighbors": 16,
+ "atom_features": "cgcnn",
+ "use_canonize": True,
+ "reduce_cell": False,
+}
+
+
+@dataclass(frozen=True)
+class _CachedPayload:
+ """Validated immutable-by-convention normalized data cache entry."""
+
+ payload: Mapping[str, Any]
+ sha256: str
+
+
+def _sha256(path: Path) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _split_indices_sha256(indices: list[int]) -> str:
+ encoded = json.dumps(indices, separators=(",", ":"), ensure_ascii=True).encode(
+ "utf-8"
+ )
+ return hashlib.sha256(encoded).hexdigest()
+
+
+def _require_regular_file(path: Path, label: str) -> Path:
+ resolved = path.expanduser().resolve()
+ if not resolved.exists():
+ raise FileNotFoundError(f"{label} does not exist: {resolved}")
+ if not resolved.is_file():
+ raise ValueError(f"{label} must be a regular file: {resolved}")
+ return resolved
+
+
+def _default_split_resource():
+ """Return the canonical split from installed resources or the source tree."""
+ try:
+ resource = resources.files("property_prediction").joinpath(
+ *_DEFAULT_SPLIT_RESOURCE
+ )
+ except ModuleNotFoundError:
+ resource = Path(__file__).resolve().parents[2].joinpath(
+ "property_prediction", *_DEFAULT_SPLIT_RESOURCE
+ )
+ if not resource.is_file():
+ raise FileNotFoundError(
+ "GMTNet split resource is missing: property_prediction/configs/gmtnet/"
+ "split_gmtnet_dielectric_seed32.json"
+ )
+ return resource
+
+
+def _validate_record(record: Any, index: int) -> None:
+ if not isinstance(record, Mapping):
+ raise ValueError(f"Record {index} must be a mapping.")
+ required_fields = {
+ "data_index",
+ "JARVIS_ID",
+ "structure",
+ "equivalent_atoms",
+ "feature_mask",
+ "matrix_equal",
+ "dielectric",
+ }
+ missing = required_fields - set(record)
+ if missing:
+ raise ValueError(f"Record {index} is missing fields: {sorted(missing)}")
+ if record["data_index"] != index:
+ raise ValueError(f"Record {index} data_index does not match its position.")
+ if not isinstance(record["JARVIS_ID"], str) or not record["JARVIS_ID"]:
+ raise ValueError(f"Record {index} JARVIS_ID must be a non-empty string.")
+ if type(record["structure"]) is not dict:
+ raise ValueError(f"Record {index} structure must be a plain dict.")
+ field_shapes = {
+ "equivalent_atoms": (None,),
+ "feature_mask": (32, 32),
+ "matrix_equal": (9, 9),
+ "dielectric": (3, 3),
+ }
+ for field_name, expected_shape in field_shapes.items():
+ value = record[field_name]
+ if not isinstance(value, np.ndarray):
+ raise ValueError(f"Record {index} {field_name} must be a numpy array.")
+ if expected_shape == (None,):
+ if value.ndim != 1:
+ raise ValueError(
+ f"Record {index} equivalent_atoms must have shape [N]."
+ )
+ elif value.shape != expected_shape:
+ raise ValueError(
+ f"Record {index} {field_name} has shape {value.shape}, expected {expected_shape}."
+ )
+
+
+def _validate_payload(payload: Any) -> Mapping[str, Any]:
+ if not isinstance(payload, Mapping):
+ raise ValueError("Normalized pickle root must be a mapping.")
+ if payload.get("schema_version") != 1:
+ raise ValueError("Normalized pickle schema_version must be 1.")
+ if payload.get("num_records") != _RECORD_COUNT:
+ raise ValueError(f"Normalized pickle num_records must be {_RECORD_COUNT}.")
+ source_hash = payload.get("source_original_dataset_sha256")
+ if not isinstance(source_hash, str) or not source_hash:
+ raise ValueError("Normalized pickle source_original_dataset_sha256 is missing.")
+ records = payload.get("records")
+ if not isinstance(records, list) or len(records) != _RECORD_COUNT:
+ raise ValueError(
+ f"Normalized pickle records must contain {_RECORD_COUNT} entries."
+ )
+ for index, record in enumerate(records):
+ _validate_record(record, index)
+ return payload
+
+
+def _validate_split_json(
+ split_path: Path,
+ normalized_sha256: str,
+ normalized_payload: Mapping[str, Any],
+ verify_sha256: bool,
+) -> dict[str, tuple[int, ...]]:
+ with split_path.open("r", encoding="utf-8") as handle:
+ split_data = json.load(handle)
+ required_fields = {
+ "schema_version",
+ "normalized_schema_version",
+ "source_original_dataset_sha256",
+ "normalized_dataset_sha256",
+ "num_records",
+ "seed",
+ "split_sizes",
+ "split_indices_sha256",
+ "train_indices",
+ "val_indices",
+ "test_indices",
+ }
+ if not isinstance(split_data, dict) or not required_fields <= set(split_data):
+ raise ValueError("Split JSON has an invalid schema.")
+ if split_data["schema_version"] != 1:
+ raise ValueError("Split JSON schema_version must be 1.")
+ if split_data["normalized_schema_version"] != 1:
+ raise ValueError("Split JSON normalized_schema_version must be 1.")
+ if split_data["num_records"] != _RECORD_COUNT:
+ raise ValueError(f"Split JSON num_records must be {_RECORD_COUNT}.")
+ if split_data["seed"] != 32:
+ raise ValueError("Split JSON seed must be 32.")
+ if split_data["split_sizes"] != _SPLIT_SIZES:
+ raise ValueError("Split JSON split_sizes are invalid.")
+ if verify_sha256 and split_data["normalized_dataset_sha256"] != normalized_sha256:
+ raise ValueError(
+ "Split JSON normalized_dataset_sha256 does not match data_path."
+ )
+ if (
+ split_data["source_original_dataset_sha256"]
+ != normalized_payload["source_original_dataset_sha256"]
+ ):
+ raise ValueError(
+ "Split JSON source_original_dataset_sha256 does not match normalized data."
+ )
+ hashes = split_data["split_indices_sha256"]
+ if not isinstance(hashes, dict) or set(hashes) != set(_SPLIT_SIZES):
+ raise ValueError("Split JSON split_indices_sha256 is invalid.")
+ split_indices: dict[str, tuple[int, ...]] = {}
+ all_indices: list[int] = []
+ for split_name, expected_size in _SPLIT_SIZES.items():
+ indices = split_data[f"{split_name}_indices"]
+ if not isinstance(indices, list) or len(indices) != expected_size:
+ raise ValueError(f"Split JSON {split_name}_indices length is invalid.")
+ if any(type(index) is not int for index in indices):
+ raise ValueError(f"Split JSON {split_name}_indices must contain integers.")
+ if any(index < 0 or index >= _RECORD_COUNT for index in indices):
+ raise ValueError(
+ f"Split JSON {split_name}_indices contains an out-of-range index."
+ )
+ if len(set(indices)) != len(indices):
+ raise ValueError(
+ f"Split JSON {split_name}_indices contains duplicate indices."
+ )
+ if hashes[split_name] != _split_indices_sha256(indices):
+ raise ValueError(f"Split JSON {split_name}_indices SHA256 does not match.")
+ split_indices[split_name] = tuple(indices)
+ all_indices.extend(indices)
+ if len(set(all_indices)) != _RECORD_COUNT or set(all_indices) != set(
+ range(_RECORD_COUNT)
+ ):
+ raise ValueError(
+ "Split JSON indices overlap or do not cover the normalized data."
+ )
+ if split_indices["test"][:3] != (747, 1423, 1322):
+ raise ValueError("Split JSON test_indices fixed prefix is invalid.")
+ return split_indices
+
+
+def _validate_smoke_split_json(
+ split_path: Path,
+ canonical_split_path: Path,
+ normalized_sha256: str,
+ normalized_payload: Mapping[str, Any],
+ verify_sha256: bool,
+) -> dict[str, tuple[int, ...]]:
+ canonical_indices = _validate_split_json(
+ canonical_split_path,
+ normalized_sha256,
+ normalized_payload,
+ verify_sha256,
+ )
+ with canonical_split_path.open("r", encoding="utf-8") as handle:
+ canonical_data = json.load(handle)
+ with split_path.open("r", encoding="utf-8") as handle:
+ split_data = json.load(handle)
+ if not isinstance(split_data, dict):
+ raise ValueError("Smoke split JSON must be a mapping.")
+ if split_data.get("smoke_only") is not True:
+ raise ValueError("Smoke split JSON smoke_only must be true.")
+ if split_data.get("parent_split_sha256") != _sha256(canonical_split_path):
+ raise ValueError(
+ "Smoke split JSON parent_split_sha256 does not match canonical split."
+ )
+ for field_name in (
+ "schema_version",
+ "normalized_schema_version",
+ "source_original_dataset_sha256",
+ "normalized_dataset_sha256",
+ "num_records",
+ "seed",
+ ):
+ if split_data.get(field_name) != canonical_data.get(field_name):
+ raise ValueError(
+ f"Smoke split JSON {field_name} does not match canonical split."
+ )
+ if verify_sha256 and split_data["normalized_dataset_sha256"] != normalized_sha256:
+ raise ValueError(
+ "Smoke split JSON normalized_dataset_sha256 does not match data_path."
+ )
+ hashes = split_data.get("split_indices_sha256")
+ if not isinstance(hashes, dict) or set(hashes) != set(_SPLIT_SIZES):
+ raise ValueError("Smoke split JSON split_indices_sha256 is invalid.")
+ split_sizes = split_data.get("split_sizes")
+ if not isinstance(split_sizes, dict) or set(split_sizes) != set(_SPLIT_SIZES):
+ raise ValueError("Smoke split JSON split_sizes is invalid.")
+ smoke_indices: dict[str, tuple[int, ...]] = {}
+ all_indices: list[int] = []
+ for split_name in _SPLIT_SIZES:
+ indices = split_data.get(f"{split_name}_indices")
+ if not isinstance(indices, list) or not indices:
+ raise ValueError(
+ f"Smoke split JSON {split_name}_indices must be a non-empty list."
+ )
+ if any(type(index) is not int for index in indices):
+ raise ValueError(
+ f"Smoke split JSON {split_name}_indices must contain integers."
+ )
+ if any(index < 0 or index >= _RECORD_COUNT for index in indices):
+ raise ValueError(
+ f"Smoke split JSON {split_name}_indices contains an out-of-range index."
+ )
+ if split_sizes[split_name] != len(indices):
+ raise ValueError(
+ f"Smoke split JSON split_sizes[{split_name}] does not match indices."
+ )
+ if hashes[split_name] != _split_indices_sha256(indices):
+ raise ValueError(
+ f"Smoke split JSON {split_name}_indices SHA256 does not match."
+ )
+ if (
+ len(indices) > len(canonical_indices[split_name])
+ or tuple(indices) != canonical_indices[split_name][: len(indices)]
+ ):
+ raise ValueError(
+ f"Smoke split JSON {split_name}_indices must be a canonical ordered prefix."
+ )
+ smoke_indices[split_name] = tuple(indices)
+ all_indices.extend(indices)
+ if len(set(all_indices)) != len(all_indices):
+ raise ValueError("Smoke split JSON indices overlap across splits.")
+ return smoke_indices
+
+
+def _converter_params(build_graph_cfg: Mapping[str, Any] | None) -> dict[str, Any]:
+ if build_graph_cfg is None:
+ return dict(_CONVERTER_DEFAULTS)
+ config = dict(build_graph_cfg)
+ if "__class_name__" in config or "__init_params__" in config:
+ class_name = config.pop("__class_name__", None)
+ init_params = config.pop("__init_params__", None)
+ if config or class_name not in {
+ "GMTNetGraphConverter",
+ "ppmat.models.gmtnet.gmtnet_graph_converter.GMTNetGraphConverter",
+ }:
+ raise ValueError("build_graph_cfg must describe GMTNetGraphConverter.")
+ if not isinstance(init_params, Mapping):
+ raise ValueError("build_graph_cfg __init_params__ must be a mapping.")
+ config = dict(init_params)
+ unsupported = set(config) - set(_CONVERTER_DEFAULTS)
+ if unsupported:
+ raise ValueError(
+ f"build_graph_cfg has unsupported GMTNetGraphConverter parameters: {sorted(unsupported)}"
+ )
+ return {**_CONVERTER_DEFAULTS, **config}
+
+
+class GMTNetDielectricDataset(Dataset):
+ """Construct frozen GMTNet dielectric samples for one fixed split."""
+
+ property_names: ClassVar[tuple[str, ...]] = ("dielectric",)
+ _payload_cache: ClassVar[dict[tuple[Path, int, int], _CachedPayload]] = {}
+ _pickle_load_count: ClassVar[int] = 0
+
+ def __init__(
+ self,
+ data_path: str | Path,
+ split: str,
+ build_graph_cfg: Mapping[str, Any] | None = None,
+ split_path: str | Path | None = None,
+ verify_sha256: bool = True,
+ allow_smoke_split: bool = False,
+ canonical_split_path: str | Path | None = None,
+ ) -> None:
+ super().__init__()
+ if not isinstance(allow_smoke_split, bool):
+ raise TypeError("allow_smoke_split must be a bool.")
+ if split not in _SPLIT_SIZES:
+ raise ValueError("split must be one of: train, val, test.")
+ self.data_path = _require_regular_file(Path(data_path), "data_path")
+ split_path_context = (
+ nullcontext(Path(split_path))
+ if split_path is not None
+ else resources.as_file(_default_split_resource())
+ )
+ with split_path_context as resolved_split_path:
+ self.split_path = _require_regular_file(
+ Path(resolved_split_path), "split_path"
+ )
+ cache_entry = self._load_payload(self.data_path)
+ self._payload = cache_entry.payload
+ if allow_smoke_split:
+ if canonical_split_path is None:
+ raise ValueError(
+ "canonical_split_path is required when allow_smoke_split is true."
+ )
+ try:
+ resolved_canonical_path = _require_regular_file(
+ Path(canonical_split_path), "canonical_split_path"
+ )
+ except (FileNotFoundError, ValueError) as error:
+ raise ValueError(
+ f"canonical_split_path is invalid: {error}"
+ ) from error
+ self._split_indices = _validate_smoke_split_json(
+ self.split_path,
+ resolved_canonical_path,
+ cache_entry.sha256,
+ self._payload,
+ verify_sha256,
+ )[split]
+ else:
+ with self.split_path.open("r", encoding="utf-8") as handle:
+ split_data = json.load(handle)
+ if isinstance(split_data, dict) and split_data.get("smoke_only") is True:
+ raise ValueError(
+ "Smoke split JSON requires allow_smoke_split to be true."
+ )
+ self._split_indices = _validate_split_json(
+ self.split_path,
+ cache_entry.sha256,
+ self._payload,
+ verify_sha256,
+ )[split]
+ self.split = split
+ self.graph_converter = GMTNetGraphConverter(
+ **_converter_params(build_graph_cfg)
+ )
+
+ @classmethod
+ def _load_payload(cls, data_path: Path) -> _CachedPayload:
+ stat = data_path.stat()
+ cache_key = (data_path, stat.st_size, stat.st_mtime_ns)
+ cached = cls._payload_cache.get(cache_key)
+ if cached is not None:
+ return cached
+ sha256 = _sha256(data_path)
+ with data_path.open("rb") as handle:
+ payload = _validate_payload(pickle.load(handle))
+ cls._pickle_load_count += 1
+ cache_entry = _CachedPayload(payload=payload, sha256=sha256)
+ cls._payload_cache[cache_key] = cache_entry
+ return cache_entry
+
+ @classmethod
+ def _clear_cache_for_testing(cls) -> None:
+ cls._payload_cache.clear()
+ cls._pickle_load_count = 0
+
+ def __len__(self) -> int:
+ return len(self._split_indices)
+
+ def __getitem__(self, index: int) -> dict[str, Any]:
+ local_index = operator.index(index)
+ if local_index < 0:
+ local_index += len(self)
+ if local_index < 0 or local_index >= len(self):
+ raise IndexError("GMTNetDielectricDataset index out of range.")
+ data_index = self._split_indices[local_index]
+ record = self._payload["records"][data_index]
+ structure = Structure.from_dict(record["structure"])
+ graph = self.graph_converter(structure, record["equivalent_atoms"])
+ return {
+ "graph": graph,
+ "feature_mask": paddle.to_tensor(record["feature_mask"], dtype="float32"),
+ "matrix_equal": paddle.to_tensor(record["matrix_equal"], dtype="bool"),
+ "dielectric": paddle.to_tensor(record["dielectric"], dtype="float32"),
+ "id": record["JARVIS_ID"],
+ "data_index": paddle.to_tensor(record["data_index"], dtype="int64"),
+ }
diff --git a/ppmat/datasets/graph_utils/spherenet_graph_utils.py b/ppmat/datasets/graph_utils/spherenet_graph_utils.py
new file mode 100644
index 00000000..c4ee4769
--- /dev/null
+++ b/ppmat/datasets/graph_utils/spherenet_graph_utils.py
@@ -0,0 +1,63 @@
+# 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.
+"""SphereNet-specific graph utilities."""
+
+import paddle
+
+
+def radius_graph(pos, batch, cutoff, loop=False):
+ """Build edge indices for a batch of molecules within a cutoff radius.
+
+ Processes each molecule independently to avoid O(N²) memory on the full
+ concatenated batch. For each molecule, builds a local N_mol × N_mol
+ distance matrix, then remaps edge indices to global positions.
+
+ Args:
+ pos: Tensor of shape ``(num_nodes, 3)`` with atomic coordinates.
+ batch: Tensor of shape ``(num_nodes,)`` with batch indices.
+ cutoff: Neighbor cutoff distance in Ångström.
+ loop: Whether to include self-loops (default False).
+
+ Returns:
+ edge_index: Tensor of shape ``(2, num_edges)`` with global edge indices.
+ """
+ num_nodes = pos.shape[0]
+ if num_nodes == 0:
+ return paddle.empty([2, 0], dtype="int64")
+
+ if batch is None:
+ batch = paddle.zeros([num_nodes], dtype="int64")
+
+ _, counts = paddle.unique(batch, return_counts=True)
+ edge_list = []
+ start = 0
+ for i in range(counts.shape[0]):
+ n = int(counts[i])
+ if n == 0:
+ continue
+ local_pos = pos[start : start + n]
+ diff = local_pos.unsqueeze(1) - local_pos.unsqueeze(0)
+ dist_sq = paddle.sum(diff * diff, axis=-1)
+ mask = dist_sq < cutoff * cutoff
+ if not loop:
+ atom_ids = paddle.arange(n, dtype="int64")
+ mask = mask & (atom_ids.unsqueeze(0) != atom_ids.unsqueeze(1))
+ src, dst = paddle.where(mask)
+ if src.shape[0] > 0:
+ edge_list.append(paddle.stack([src + start, dst + start], axis=0))
+ start += n
+
+ if not edge_list:
+ return paddle.empty([2, 0], dtype="int64")
+ return paddle.concat(edge_list, axis=1)
diff --git a/ppmat/datasets/md17_dataset.py b/ppmat/datasets/md17_dataset.py
new file mode 100644
index 00000000..c5bd4fa7
--- /dev/null
+++ b/ppmat/datasets/md17_dataset.py
@@ -0,0 +1,392 @@
+# 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 molecular dynamics dataset for energy and force prediction."""
+
+import os
+import os.path as osp
+import pickle
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Optional
+
+import numpy as np
+import paddle.distributed as dist
+from paddle.io import Dataset
+
+from ppmat.datasets.build_molecule import BuildMolecule
+from ppmat.datasets.custom_data_type import ConcatData
+from ppmat.models import build_graph_converter
+from ppmat.utils import download
+from ppmat.utils import logger
+from ppmat.utils.misc import is_equal
+
+
+_BUNDLE_NPZ_MAP = {
+ "aspirin": "md17_aspirin.npz",
+ "benzene_old": "md17_benzene2017.npz",
+ "ethanol": "md17_ethanol.npz",
+ "malonaldehyde": "md17_malonaldehyde.npz",
+ "naphthalene": "md17_naphthalene.npz",
+ "salicylic": "md17_salicylic.npz",
+ "toluene": "md17_toluene.npz",
+ "uracil": "md17_uracil.npz",
+}
+
+
+class MD17Dataset(Dataset):
+ """MD17 molecular dynamics dataset for energy and force prediction.
+
+ **STATS:**
+ +----------------+----------+--------+-------+----------+-------+
+ | Molecule | #samples | #atoms | #tasks| #targets | Split |
+ +================+==========+========+=======+==========+=======+
+ | Aspirin | 211,762 | 21 | 2 | E + F | 1k/1k/R |
+ | Benzene (old) | 627,983 | 12 | 2 | E + F | 1k/1k/R |
+ | Ethanol | 555,092 | 9 | 2 | E + F | 1k/1k/R |
+ | Malonaldehyde | 993,237 | 9 | 2 | E + F | 1k/1k/R |
+ | Naphthalene | 326,250 | 10 | 2 | E + F | 1k/1k/R |
+ | Salicylic | 320,231 | 16 | 2 | E + F | 1k/1k/R |
+ | Toluene | 442,790 | 15 | 2 | E + F | 1k/1k/R |
+ | Uracil | 133,770 | 12 | 2 | E + F | 1k/1k/R |
+ +----------------+----------+--------+-------+----------+-------+
+
+ Contains ab-initio molecular dynamics trajectories for eight small
+ organic molecules. Each frame provides atomic numbers, 3D positions,
+ total energy, and per-atom forces.
+
+ Data source: https://www.quantum-machine.org/datasets/
+
+ Args:
+ path (str): Root directory for storing raw and cached data.
+ name (str): Molecule name from the supported list. Defaults to ``'benzene_old'``.
+ split (Optional[str]): Split identifier ``'train'``, ``'val'``,
+ ``'test'``, or ``None`` (all). Defaults to ``None``.
+ split_file (Optional[str]): Preprocessed numpy index file for the
+ selected split. Defaults to ``None``.
+ build_molecule_cfg (Optional[Dict]): Configuration dict for molecule
+ converter. Defaults to ``None``.
+ force_key (Optional[str]): Key name for forces in the output dict.
+ Defaults to ``'force'``.
+ energy_key (Optional[str]): Key name for energy in the output dict.
+ Defaults to ``'energy'``.
+ build_graph_cfg (Optional[Dict]): Configuration dict for graph
+ converter. Defaults to ``None``.
+ transforms (Optional[Callable]): Per-sample transform callable.
+ Defaults to ``None``.
+ cache_path (Optional[str]): Explicit cache path. Auto-generated
+ when ``None``. Defaults to ``None``.
+ overwrite (bool): Whether to overwrite existing cached graphs.
+ Defaults to ``False``.
+ filter_unvalid (bool): Whether to filter out invalid samples.
+ Defaults to ``True``.
+ """
+
+ url = "https://paddle-org.bj.bcebos.com/paddlematerials/datasets/MD17/md17.tar.gz"
+ md5 = "634cc25cc8a3fb0d99bd14245eb8dabd"
+ name = "md17"
+
+ def __init__(
+ self,
+ path: str,
+ name: str = "benzene_old",
+ split: str = None,
+ *,
+ split_file: Optional[str] = None,
+ build_molecule_cfg: Optional[Dict] = None,
+ force_key="force",
+ energy_key="energy",
+ build_graph_cfg: Optional[Dict] = None,
+ transforms: Optional[Callable] = None,
+ cache_path: Optional[str] = None,
+ overwrite: bool = False,
+ filter_unvalid: bool = True,
+ **kwargs,
+ ):
+ super().__init__()
+
+ self.mol_name = name
+ self.force_key = force_key
+ self.energy_key = energy_key
+ self.transforms = transforms
+ self.overwrite = overwrite
+ self.filter_unvalid = filter_unvalid
+ if build_molecule_cfg is None:
+ build_molecule_cfg = {
+ "format": "dict",
+ "sanitize": False,
+ "add_hs": False,
+ "remove_hs": False,
+ "kekulize": False,
+ "num_cpus": 1,
+ }
+ logger.message(
+ "The build_molecule_cfg is not set, will use the default "
+ f"configs: {build_molecule_cfg}"
+ )
+ self.build_molecule_cfg = build_molecule_cfg
+
+ # ---- 1. Read full trajectory and selected frame indices ----
+ npz_path = self._resolve_data_path(path, name)
+ self.path = npz_path
+ self.row_data, total = self.read_data(npz_path)
+ self._indices = self._load_split_indices(
+ npz_path, name, split, split_file, total
+ )
+ self.num_samples = len(self._indices)
+
+ # ---- 2. Cache path ----
+ if cache_path is not None:
+ self.cache_path = cache_path
+ else:
+ base_dir = osp.split(npz_path)[0]
+ base_name = osp.splitext(osp.basename(npz_path))[0]
+ split_suffix = split if split is not None else "all"
+ self.cache_path = osp.join(
+ f"{base_dir}_cache", f"{base_name}_{split_suffix}"
+ )
+ logger.info(f"Cache path: {self.cache_path}")
+
+ # ---- 3. Pre-build molecules and graphs (MPtrj-style cache) ----
+ self.cache_exists = True if osp.exists(self.cache_path) else False
+ if self.cache_exists and not overwrite:
+ logger.warning(
+ "Cache enabled. If a cache file exists, it will be automatically "
+ "read and current settings will be ignored. Please ensure that the "
+ "settings used in match your current settings."
+ )
+ try:
+ build_molecule_cfg_cache = self.load_from_cache(
+ osp.join(self.cache_path, "build_molecule_cfg.pkl")
+ )
+ if not is_equal(build_molecule_cfg_cache, build_molecule_cfg):
+ logger.warning(
+ "build_molecule_cfg differs from cache. Rebuilding."
+ )
+ overwrite = True
+ except Exception as e:
+ logger.warning(e)
+ logger.warning(
+ "Failed to load build_molecule_cfg.pkl from cache. "
+ "Will rebuild the molecules and graphs(if need)."
+ )
+ overwrite = True
+
+ if build_graph_cfg is not None and not overwrite:
+ try:
+ build_graph_cfg_cache = self.load_from_cache(
+ osp.join(self.cache_path, "build_graph_cfg.pkl")
+ )
+ if not is_equal(build_graph_cfg_cache, build_graph_cfg):
+ logger.warning(
+ "build_graph_cfg differs from cache. Rebuilding."
+ )
+ overwrite = True
+ except Exception as e:
+ logger.warning(e)
+ logger.warning(
+ "Failed to load build_graph_cfg.pkl from cache. "
+ "Will rebuild the graphs."
+ )
+ overwrite = True
+
+ molecule_cache_path = osp.join(self.cache_path, "molecules")
+ graph_cache_path = osp.join(self.cache_path, "graphs")
+ if overwrite or not self.cache_exists:
+ if dist.get_rank() == 0:
+ os.makedirs(self.cache_path, exist_ok=True)
+ self.save_to_cache(
+ osp.join(self.cache_path, "build_molecule_cfg.pkl"),
+ build_molecule_cfg,
+ )
+ self.save_to_cache(
+ osp.join(self.cache_path, "build_graph_cfg.pkl"), build_graph_cfg
+ )
+
+ molecule_data = [
+ {
+ "atomic_numbers": self.row_data["z"],
+ "positions": self.row_data["pos"][frame],
+ }
+ for frame in self._indices
+ ]
+ molecules = BuildMolecule(**build_molecule_cfg)(molecule_data)
+ os.makedirs(molecule_cache_path, exist_ok=True)
+ for i, mol in enumerate(molecules):
+ self.save_to_cache(
+ osp.join(molecule_cache_path, f"{i:010d}.pkl"), mol
+ )
+ logger.info(
+ f"Save {self.num_samples} molecules to {molecule_cache_path}"
+ )
+
+ if build_graph_cfg is not None:
+ converter = build_graph_converter(build_graph_cfg)
+ graphs = converter(molecules)
+ os.makedirs(graph_cache_path, exist_ok=True)
+ for i in range(self.num_samples):
+ self.save_to_cache(
+ osp.join(graph_cache_path, f"{i:010d}.pkl"), graphs[i]
+ )
+ logger.info(f"Save {self.num_samples} graphs to {graph_cache_path}")
+ if dist.is_initialized():
+ dist.barrier()
+
+ self.molecules = [
+ osp.join(molecule_cache_path, f"{i:010d}.pkl")
+ for i in range(self.num_samples)
+ ]
+ if build_graph_cfg is not None:
+ self.graphs = [
+ osp.join(graph_cache_path, f"{i:010d}.pkl")
+ for i in range(self.num_samples)
+ ]
+ else:
+ self.graphs = None
+ assert (
+ len(self.molecules) == self.num_samples
+ ), "The number of molecules must be equal to the number of samples."
+ assert (
+ self.graphs is None or len(self.graphs) == self.num_samples
+ ), "The number of graphs must be equal to the number of samples."
+
+ logger.info(f"Load {self.num_samples} samples, split={split}")
+
+ def _resolve_data_path(self, path, name):
+ if osp.isfile(path):
+ return path
+
+ candidates = [
+ osp.join(path, f"{name}_dft.npz"),
+ osp.join(path, self.name, f"{name}_dft.npz"),
+ ]
+ if name in _BUNDLE_NPZ_MAP:
+ candidates.extend(
+ [
+ osp.join(path, _BUNDLE_NPZ_MAP[name]),
+ osp.join(path, self.name, _BUNDLE_NPZ_MAP[name]),
+ ]
+ )
+ for candidate in candidates:
+ if osp.exists(candidate):
+ return candidate
+
+ logger.message("The dataset is not found. Will download it now.")
+ root_path = download.get_datasets_path_from_url(self.url, self.md5)
+ root_paths = [root_path]
+ if not osp.exists(root_path):
+ root_paths.append(osp.dirname(root_path))
+
+ candidates = []
+ for candidate_root in root_paths:
+ candidates.extend(
+ [
+ osp.join(candidate_root, self.name, f"{name}_dft.npz"),
+ osp.join(candidate_root, f"{name}_dft.npz"),
+ ]
+ )
+ if name in _BUNDLE_NPZ_MAP:
+ candidates.extend(
+ [
+ osp.join(
+ candidate_root, self.name, _BUNDLE_NPZ_MAP[name]
+ ),
+ osp.join(candidate_root, _BUNDLE_NPZ_MAP[name]),
+ ]
+ )
+ for candidate in candidates:
+ if osp.exists(candidate):
+ return candidate
+ raise FileNotFoundError(f"Cannot find MD17 npz file for molecule: {name}")
+
+ def read_data(self, path):
+ """Load all trajectory frames from the npz file."""
+ data = np.load(path)
+ row_data = {
+ "z": data["z"],
+ "pos": data["R"],
+ "energy": data["E"],
+ "force": data["F"],
+ }
+ return row_data, data["R"].shape[0]
+
+ def _load_split_indices(self, path, name, split, split_file, total):
+ """Load frame indices for the requested split."""
+ if split is None and split_file is None:
+ return np.arange(total, dtype=np.int64)
+ if split_file is None:
+ split_dir = osp.join(osp.dirname(path), "splits")
+ key = split if split is not None else "all"
+ split_file = osp.join(split_dir, f"{name}_{key}_idx.npy")
+ if not osp.exists(split_file):
+ split_file = osp.join(split_dir, f"split_{key}.npy")
+ if not osp.exists(split_file):
+ raise FileNotFoundError(f"No such split file: {split_file}")
+ return np.load(split_file).astype(np.int64)
+
+ def get_molecule_array(self, molecule):
+ conf = molecule.GetConformer()
+ z = np.array(
+ [atom.GetAtomicNum() for atom in molecule.GetAtoms()], dtype=np.int64
+ )
+ pos = np.array(
+ [
+ [
+ conf.GetAtomPosition(i).x,
+ conf.GetAtomPosition(i).y,
+ conf.GetAtomPosition(i).z,
+ ]
+ for i in range(molecule.GetNumAtoms())
+ ],
+ dtype=np.float32,
+ )
+ return {
+ "z": ConcatData(z),
+ "pos": ConcatData(pos),
+ "num_atoms": ConcatData(np.array([z.shape[0]], dtype=np.int64)),
+ }
+
+ def save_to_cache(self, cache_path: str, obj: Any):
+ with open(cache_path, "wb") as f:
+ pickle.dump(obj, f)
+
+ def load_from_cache(self, cache_path: str):
+ if osp.exists(cache_path):
+ with open(cache_path, "rb") as f:
+ return pickle.load(f)
+ raise FileNotFoundError(f"No such file or directory: {cache_path}")
+
+ def __getitem__(self, idx):
+ frame = self._indices[idx]
+ sample = {
+ self.energy_key: np.array(
+ [float(self.row_data["energy"][frame])], dtype=np.float32
+ ),
+ self.force_key: ConcatData(self.row_data["force"][frame]),
+ }
+ if self.graphs is not None:
+ graph = self.graphs[idx]
+ if isinstance(graph, str):
+ graph = self.load_from_cache(graph)
+ sample["graph"] = graph
+ else:
+ mol = self.load_from_cache(self.molecules[idx])
+ sample.update(self.get_molecule_array(mol))
+ sample["id"] = int(frame)
+ if self.transforms is not None:
+ sample = self.transforms(sample)
+ return sample
+
+ def __len__(self):
+ return self.num_samples
diff --git a/ppmat/datasets/mp2018_dataset.py b/ppmat/datasets/mp2018_dataset.py
index 2c0fe2a8..ea4518e8 100644
--- a/ppmat/datasets/mp2018_dataset.py
+++ b/ppmat/datasets/mp2018_dataset.py
@@ -185,9 +185,9 @@ def __init__(
self.filter_unvalid = filter_unvalid
self.cache_exists = True if osp.exists(self.cache_path) else False
- self.row_data, self.num_samples = self.read_data(path)
+ self.raw_data, self.num_samples = self.read_data(path)
logger.info(f"Load {self.num_samples} samples from {path}")
- self.property_data = self.read_property_data(self.row_data, self.property_names)
+ self.property_data = self.read_property_data(self.raw_data, self.property_names)
if self.cache_exists and not overwrite:
logger.warning(
@@ -272,7 +272,7 @@ def __init__(
)
# convert strucutes
structures = BuildStructure(**build_structure_cfg)(
- self.row_data["structure"]
+ self.raw_data["structure"]
)
# save structures to cache file
os.makedirs(structure_cache_path, exist_ok=True)
@@ -352,7 +352,7 @@ def filter_unvalid_by_property(self):
self.property_data[key][i] for i in reserve_idx
]
- self.row_data = [self.row_data[i] for i in reserve_idx]
+ self.raw_data = [self.raw_data[i] for i in reserve_idx]
self.structures = [self.structures[i] for i in reserve_idx]
if self.graphs is not None:
self.graphs = [self.graphs[i] for i in reserve_idx]
@@ -360,7 +360,7 @@ def filter_unvalid_by_property(self):
f"Filter out {len(reserve_idx)} samples with valid properties: "
f"{property_name}"
)
- self.num_samples = len(self.row_data)
+ self.num_samples = len(self.raw_data)
logger.warning(f"Remaining {self.num_samples} samples after filtering.")
def read_property_data(self, data: Dict, property_names: list[str]):
diff --git a/ppmat/datasets/mp20_dataset.py b/ppmat/datasets/mp20_dataset.py
index d257f855..a70e778d 100644
--- a/ppmat/datasets/mp20_dataset.py
+++ b/ppmat/datasets/mp20_dataset.py
@@ -128,6 +128,7 @@ def __init__(
self.path = path
if isinstance(property_names, str):
property_names = [property_names]
+ self.property_names = property_names if property_names is not None else []
if build_structure_cfg is None:
build_structure_cfg = {
@@ -141,7 +142,6 @@ def __init__(
f"configs: {build_structure_cfg}"
)
- self.property_names = property_names if property_names is not None else []
self.build_structure_cfg = build_structure_cfg
self.build_graph_cfg = build_graph_cfg
self.transforms = transforms
diff --git a/ppmat/datasets/msd_nmr_dataset.py b/ppmat/datasets/msd_nmr_dataset.py
index f3e08798..2406bb5e 100644
--- a/ppmat/datasets/msd_nmr_dataset.py
+++ b/ppmat/datasets/msd_nmr_dataset.py
@@ -277,12 +277,12 @@ def __init__(
logger.info(
"The cached build_molecule_cfg configuration matches "
"the current settings. Reusing previously generated"
- " structural data to optimize performance."
+ " moelcular data to optimize performance."
)
else:
logger.warning(
"build_molecule_cfg is different from "
- "build_molecule_cfgg_cache. Will rebuild the molecules and "
+ "build_molecule_cfg_cache. Will rebuild the molecules and "
"graphs."
)
logger.warning(
@@ -365,7 +365,7 @@ def __init__(
spectrum_cache_path = osp.join(self.cache_path, "spectrums")
if overwrite or not self.cache_exists:
- # convert strucutes and graphs
+ # convert molecules and graphs
# only rank 0 process do the conversion
if dist.get_rank() == 0:
# save build_molecule_cfg and build_graph_cfg and build_spechtrum_cfg
diff --git a/ppmat/datasets/qm9_dataset.py b/ppmat/datasets/qm9_dataset.py
index fec8889c..837049e3 100644
--- a/ppmat/datasets/qm9_dataset.py
+++ b/ppmat/datasets/qm9_dataset.py
@@ -1,766 +1,383 @@
-# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
-
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
-
+#
# http://www.apache.org/licenses/LICENSE-2.0
-
+#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
-from __future__ import absolute_import
from __future__ import annotations
import math
import os
import os.path as osp
-import pickle ## for dump/load
-from collections import defaultdict
+import pickle
from typing import Any
from typing import Callable
from typing import Dict
from typing import Optional
-from typing import Union
-from typing import List
import numpy as np
+import pandas as pd
import paddle.distributed as dist
from paddle.io import Dataset
-
-from ppmat.datasets.build_structure import BuildStructure
-from ppmat.datasets.custom_data_type import ConcatData
+from ppmat.datasets.build_molecule import BuildMolecule
from ppmat.models import build_graph_converter
from ppmat.utils import download
from ppmat.utils import logger
-from ppmat.utils.io import read_json
from ppmat.utils.misc import is_equal
-# Attempt to import tqdm for progress visualization
-try:
- from tqdm import tqdm
-except ImportError:
- def tqdm(iterable, **kwargs):
- return iterable
-
-try:
- import ase.io
- import ase.data
- import ase.build
- from pymatgen.io.ase import AseAtomsAdaptor
-
- read = ase.io.read
- symbols = ase.data.atomic_numbers
- ASE_AVAILABLE = True
-
-except ImportError:
-
- def dummy_read(*args, **kwargs):
- """
- If ASE is not installed, this pseudo-read function will throw an error when called.
- In the `__init__` method of `QM9Dataset`, if `read` is not needed immediately,
- you can simply return `None` or an empty list. However, if it is needed, the error will be more explicit.
-
- """
- raise RuntimeError(
- "Atomic Simulation Environment (ASE) is required but not installed. "
- "Please install it (e.g., pip install ase) to use QM9Dataset."
- )
- read = dummy_read
- symbols = None
- AseAtomsAdaptor = None
- ASE_AVAILABLE = False
- print("Warning: ASE (Atomic Simulation Environment) not found. Data parsing functionality is disabled.")
class QM9Dataset(Dataset):
- """
- QM9 (GDB-9) Dataset Handler
- In order to adapt to the CHGNet model, I made a forced mapping from 'lumo' to 'energy_per_atom'.
- The CHGNet model is used to run machine learning potentials, and the QM9 dataset may not be very suitable.
- LUMO is a property. Please be aware of this when using it to avoid misunderstandings.
- This code is for research purposes only and does not represent the optimal approach.——by wwaawwaaee
-
+ """QM9 Dataset Handler
+
+ This class provides utilities for loading and processing the QM9 molecular
+ quantum chemistry dataset. The implementation supports both standard dataset
+ loading and custom data processing when adhering to the QM9 data schema.
+
**Dataset Overview**
- this class downloads QM9 primiry(end with .xyz)and transfered into
- the input structures and quantum chemical property labels required
- by graph neural network (GNN) models.(maybe)
-
- **dataset format**
- -----------------
- - raw data: qm9.zip ()
- - struncture file: a sample corresponds to a seperate .xyz file
- - attribute(label): 19 quantum chemical properties are embedded in each .xyz file's
- second line comment
-
- **source**:Original data available at https://figshare.com/ndownloader/files/3195389
-
- The dataset can also be found at https://paddle-org.bj.bcebos.com/paddlematerials/datasets/qm9/dsgdb9nsd.xyz.tar.bz2
-
- **Key Properties List (Available for 'property_names' argument)**
- -----------------------------------------------------------------
- 1. mu (Dipole Moment, Debye)
- 2. alpha (Isotropic Polarizability, Bohr^3)
- 3. homo (HOMO Energy, Hartree)
- 4. lumo (LUMO Energy, Hartree)
- 5. gap (LUMO-HOMO Gap, Hartree)
- 6. U0 (Internal Energy at 0 K, Hartree)
- sly
- # ... (remaining 13 properties here in their correct order)
-
- **__getitem__ Sample Contract**
- ----------------------------------------
- - 'atom_types': np.ndarray (dtype=int64) - Atomic numbers (Z).
- - 'coords': np.ndarray (dtype=float32) - 3D Cartesian coordinates in Angstrom.
- - [property_name]: np.ndarray (dtype=float32) - The target label value (e.g., 'lumo').
- - 'graph': (Optional) The graph object constructed by the converter (if configured).
-
+ - **Source**: Original data from "Quantum chemistry structures and properties
+ of 134 kilo molecules" and the QM9/GDB9 dataset.
+ - **Filtering**: The 3,054 uncharacterized molecules listed by the official
+ QM9 consistency check are removed.
+ ```
+ ┌───────────────────┬─────────┬─────────┬─────────┐
+ │ Dataset Partition │ Train │ Val │ Test │
+ ├───────────────────┼─────────┼─────────┼─────────┤
+ │ Sample Count │ 110,000 │ 10,000 │ 10,831 │
+ └───────────────────┴─────────┴─────────┴─────────┘
+ ```
+ The dataset can also be downloaded from the following source:
+ https://paddle-org.bj.bcebos.com/paddlematerials/datasets/qm9/qm9_split.zip
+
+ **Data Format**
+ The dataset is structured as comma-separated values (CSV) files with one
+ molecule per row. The split files are `train.csv`, `val.csv`, and `test.csv`.
+
+ | Column Name | Description | Example Value |
+ |----------------------------|--------------------------------------------------|----------------|
+ | `file_name` | Original QM9 xyz file name | qm9_000001.xyz |
+ | `standard_xyz` | Standard XYZ string without Mulliken charges | xyz_str |
+ | `mulliken_xyz` | XYZ string with Mulliken partial charges | xyz_charge_str |
+ | `molecule_id` | 1-based QM9 molecule identifier | 1 |
+ | `num_atoms` | Number of atoms in the molecule | 5 |
+ | `A`, `B`, `C` | Rotational constants | 157.7118 |
+ | `mu` | Dipole moment | 0.0 |
+ | `alpha` | Isotropic polarizability | 13.21 |
+ | `homo`, `lumo`, `gap` | HOMO, LUMO, and HOMO-LUMO gap | -0.3877 |
+ | `r2` | Electronic spatial extent | 35.3641 |
+ | `zpve` | Zero point vibrational energy | 0.044749 |
+ | `U0`, `U`, `H`, `G` | Thermochemical energies | -40.47893 |
+ | `Cv` | Heat capacity | 6.469 |
+ | `vibrational_frequencies` | Space-separated vibrational frequencies | 1341.307 ... |
+ | `canonical_smiles` | Canonical SMILES string | C |
+ | `isomeric_smiles` | Isomeric SMILES string | C |
+ | `canonical_inchi` | Canonical InChI string | InChI=1S/CH4/h1H4 |
+ | `isomeric_inchi` | Isomeric InChI string | InChI=1S/CH4/h1H4 |
+
+ **Example Row:**
+ ```csv
+ file_name,molecule_id,num_atoms,A,B,C,mu,alpha,homo,lumo,gap,...
+ qm9_000001.xyz,1,5,157.7118,157.70997,157.70699,0.0,13.21,-0.3877,0.1171,0.5048,...
+ ```
+
Args:
- path (str): The root directory to store downloaded and cache files.
- property_names (Union[str, List[str]]): The name(s) of the target property
- to predict. Must be selected from the list above. Defaults to 'lumo'.
- build_graph_cfg (Dict, optional): Configuration dictionary for building
- the graph representation from the molecular structure (e.g., cutoff radius).
- Defaults to None (structure is returned instead of graph).
- transforms (Optional[Callable], optional): A preprocessing function to apply
- to the sample dictionary. Defaults to None.
- cache_path (Optional[str], optional): Explicit path for the cache directory.
- Defaults to None.
- overwrite (bool, optional): If True, forces the rebuilding of caches.
- Defaults to False.
- filter_unvalid (bool, optional): Whether to filter out corrupted samples.
+ path (str, optional): The path of the dataset. If the path does not exist,
+ it will be downloaded. Defaults to "./data/qm9/train.csv".
+ property_names (Optional[list[str]], optional): Property names to use for
+ QM9. The property_names should be selected from
+ ["A", "B", "C", "mu", "alpha", "homo", "lumo", "gap", "r2",
+ "zpve", "U0", "U", "H", "G", "Cv"]. Defaults to None.
+ build_molecule_cfg (Dict, optional): Configs for building molecular
+ structures from xyz strings. If not specified, the default setting
+ will be used. Defaults to None.
+ build_graph_cfg (Dict, optional): Configs for building molecular graphs
+ from structures. Defaults to None.
+ transforms (Optional[Callable], optional): Preprocess transforms for each
+ sample. Defaults to None.
+ cache_path (Optional[str], optional): If a cache_path is set, parsed
+ molecules and graphs will be read directly from this path; if the
+ cache does not exist, the converted molecules and graphs will be
+ saved to this path. Defaults to None.
+ overwrite (bool, optional): Overwrite the existing cache file at the given
+ path if it already exists. Defaults to False.
+ filter_unvalid (bool, optional): Whether to filter out invalid samples.
Defaults to True.
"""
- url = "https://paddle-org.bj.bcebos.com/paddlematerials/datasets/qm9/dsgdb9nsd.xyz.tar.bz2"
name = "qm9"
- md5 = "AD1EBD51EE7F5B3A6E32E974E5D54012"
-
- # Official QM9 second-line property order (including tag/index)
- PROP_ORDER = [
- "tag", # textual tag / molecule identifier (often 'gdb ...')
- "index", # numeric index (maps to vals_float[0] after removing 'gdb')
- "A", # rotational constant A (GHz)
- "B", # rotational constant B (GHz)
- "C", # rotational constant C (GHz)
- "mu", # dipole moment (Debye)
- "alpha", # isotropic polarizability (Bohr^3)
- "homo", # HOMO energy (Hartree)
- "lumo", # LUMO energy (Hartree)
- "gap", # LUMO-HOMO gap (Hartree)
- "r2", # electronic spatial extent (Bohr^2)
- "zpve", # zero point vibrational energy (Hartree)
- "U0", # internal energy at 0K (Hartree)
- "U", # internal energy at 298.15 K (Hartree)
- "H", # enthalpy at 298.15 K (Hartree)
- "G", # free energy at 298.15 K (Hartree)
- "Cv", # heat capacity at 298.15 K (cal/mol/K)
- ]
+ url = "https://paddle-org.bj.bcebos.com/paddlematerials/datasets/qm9/qm9.zip"
+ md5 = "a70eb6cc913427db1fc1f8d3631fe00a"
def __init__(
self,
- path: str,
- url: Optional[str] = None,
- property_names: Union[str, List[str]] = None,
- *,
- url_indices: Optional[List[int]] = None,
+ path: str = "./data/qm9/train.csv",
+ property_names: Optional[list[str]] = None,
+ build_molecule_cfg: Dict = None,
build_graph_cfg: Dict = None,
transforms: Optional[Callable] = None,
cache_path: Optional[str] = None,
overwrite: bool = False,
filter_unvalid: bool = True,
**kwargs,
- ) -> None:
+ ):
super().__init__()
- # Use the ASE_AVAILABLE flag and AseAtomsAdaptor presence to validate dependencies
- if not ASE_AVAILABLE or AseAtomsAdaptor is None:
- raise RuntimeError(
- "QM9Dataset requires 'ase' and 'pymatgen'. "
- "Please install them via: pip install ase pymatgen"
- )
-
- if property_names is None:
- raise ValueError("property_names must be provided for QM9Dataset")
+ if not osp.exists(path):
+ logger.message("The dataset is not found. Will download it now.")
+ root_path = download.get_datasets_path_from_url(self.url, self.md5)
+ path = osp.join(root_path, self.name, osp.basename(path))
- if isinstance(property_names,str):
+ self.path = path
+ if isinstance(property_names, str):
property_names = [property_names]
- self.property_names = list(property_names) if property_names else []
-
- # Handle URLs configuration
- self.url = url if url is not None else self.url
+ self.property_names = property_names if property_names is not None else []
+
+ if build_molecule_cfg is None:
+ build_molecule_cfg = {
+ "format": "xyz_block",
+ "sanitize": False,
+ "add_hs": False,
+ "remove_hs": False,
+ "kekulize": False,
+ "num_cpus": 1,
+ }
+ logger.message(
+ "The build_molecule_cfg is not set, will use the default "
+ f"configs: {build_molecule_cfg}"
+ )
- #Path Configuration
- os.makedirs(path, exist_ok=True)
- self.raw_dir = osp.join(path, "raw_qm9")
- os.makedirs(self.raw_dir, exist_ok=True)
+ self.build_molecule_cfg = build_molecule_cfg
+ self.build_graph_cfg = build_graph_cfg
+ self.transforms = transforms
- self.raw_xyz_path = osp.join(self.raw_dir, "dsgdb9nsd.xyz")
-
- # Generate cache directory naming based on graph config
- if build_graph_cfg is not None:
- graph_converter_name = build_graph_cfg.get("__class_name__", "custom")
- cutoff_name = str(
- int(build_graph_cfg.get("__init_params__", {}).get("cutoff", 5))
- )
+ if cache_path is not None:
+ self.cache_path = cache_path
else:
- graph_converter_name = "none"
- cutoff_name = "none"
-
- base_cache = cache_path if cache_path is not None else path #determine the final path
- self.cache_path = osp.join(
- base_cache,
- f"qm9_cache_{graph_converter_name}_cutoff_{cutoff_name}",
- )
-
- self.transforms = transforms
+ self.cache_path = osp.join(
+ osp.split(path)[0] + "_cache", osp.splitext(osp.basename(path))[0]
+ )
+ logger.info(f"Cache path: {self.cache_path}")
+
self.overwrite = overwrite
self.filter_unvalid = filter_unvalid
- self.build_graph_cfg = build_graph_cfg
- # define sub-directories for cache
- self.structures_dir = osp.join(self.cache_path, "structures")
- self.graphs_dir = osp.join(self.cache_path, "graphs")
- self.props_dir = osp.join(self.cache_path, "properties")
-
- if dist.get_rank() == 0:
- logger.info(f"Cache path: {self.cache_path}")
- os.makedirs(self.structures_dir, exist_ok=True)
- os.makedirs(self.graphs_dir, exist_ok=True)
- os.makedirs(self.props_dir, exist_ok=True)
-
- # =========== data operation ==============
-
- # 1) Download and ensure shard files exist locally
- local_raw_file = self._ensure_raw_data()
-
- # 2) Check or build Structures and Properties cache
- if dist.get_rank() == 0:
- self._prepare_structures_and_properties(local_raw_file)
- # Only rank 0 performs the build process to avoid race conditions
-
- if dist.is_initialized():
- dist.barrier()
-
- # 3) Check or build Graphs cache (if configuration provided)
- if self.build_graph_cfg is not None:
+ self.cache_exists = True if osp.exists(self.cache_path) else False
+ self.raw_data, self.num_samples = self.read_data(path)
+ logger.info(f"Load {self.num_samples} samples from {path}")
+ self.property_data = self.read_property_data(self.raw_data, self.property_names)
+
+ if self.cache_exists and not overwrite:
+ logger.warning(
+ "Cache enabled. If a cache file exists, it will be automatically "
+ "read and current settings will be ignored. Please ensure that the "
+ "settings used in match your current settings."
+ )
+ try:
+ build_molecule_cfg_cache = self.load_from_cache(
+ osp.join(self.cache_path, "build_molecule_cfg.pkl")
+ )
+ if is_equal(build_molecule_cfg_cache, build_molecule_cfg):
+ logger.info(
+ "The cached build_molecule_cfg configuration matches "
+ "the current settings. Reusing previously generated"
+ " molecular data to optimize performance."
+ )
+ else:
+ logger.warning(
+ "build_molecule_cfg is different from "
+ "build_molecule_cfg_cache. Will rebuild the molecules and "
+ "graphs."
+ )
+ logger.warning(
+ "If you want to use the cached molecules and graphs, please "
+ "ensure that the settings used in match your current settings."
+ )
+ overwrite = True
+ except Exception as e:
+ logger.warning(e)
+ logger.warning(
+ "Failed to load builded_molecules_cfg.pkl from cache. "
+ "Will rebuild the molecules and graphs(if need)."
+ )
+ overwrite = True
+
+ if build_graph_cfg is not None and not overwrite:
+ try:
+ build_graph_cfg_cache = self.load_from_cache(
+ osp.join(self.cache_path, "build_graph_cfg.pkl")
+ )
+ if is_equal(build_graph_cfg_cache, build_graph_cfg):
+ logger.info(
+ "The cached build_molecule_cfg configuration "
+ "matches the current settings. Reusing previously "
+ "generated molecular data to optimize performance."
+ )
+ else:
+ logger.warning(
+ "build_graph_cfg is different from build_graph_cfg_cache"
+ ". Will rebuild the graphs."
+ )
+ logger.warning(
+ "If you want to use the cached molecules and graphs, "
+ "please ensure that the settings used in match your "
+ "current settings."
+ )
+ overwrite = True
+
+ except Exception as e:
+ logger.warning(e)
+ logger.warning(
+ "Failed to load builded_graph_cfg.pkl from cache. "
+ "Will rebuild the graphs."
+ )
+ overwrite = True
+
+ molecule_cache_path = osp.join(self.cache_path, "molecules")
+ graph_cache_path = osp.join(self.cache_path, "graphs")
+ if overwrite or not self.cache_exists:
if dist.get_rank() == 0:
- self._prepare_graphs()
+ os.makedirs(self.cache_path, exist_ok=True)
+ self.save_to_cache(
+ osp.join(self.cache_path, "build_molecule_cfg.pkl"),
+ build_molecule_cfg,
+ )
+ self.save_to_cache(
+ osp.join(self.cache_path, "build_graph_cfg.pkl"), build_graph_cfg
+ )
+
+ molecules = BuildMolecule(**build_molecule_cfg)(
+ self.raw_data["molecule"]
+ )
+ os.makedirs(molecule_cache_path, exist_ok=True)
+ for i in range(self.num_samples):
+ self.save_to_cache(
+ osp.join(molecule_cache_path, f"{i:010d}.pkl"), molecules[i]
+ )
+ logger.info(
+ f"Save {self.num_samples} molecules to {molecule_cache_path}"
+ )
+
+ if build_graph_cfg is not None:
+ converter = build_graph_converter(build_graph_cfg)
+ graphs = converter(molecules)
+ os.makedirs(graph_cache_path, exist_ok=True)
+ for i in range(self.num_samples):
+ self.save_to_cache(
+ osp.join(graph_cache_path, f"{i:010d}.pkl"), graphs[i]
+ )
+ logger.info(f"Save {self.num_samples} graphs to {graph_cache_path}")
+
if dist.is_initialized():
dist.barrier()
-
- PROPERTY_FILE_MAP = {
- "energy_per_atom": "lumo", # cheat the model of the way it gets the data
- }
- # 4) Load file lists and property data into memory
- self.structures = [
- osp.join(self.structures_dir, f)
- for f in sorted(os.listdir(self.structures_dir))
- if f.endswith(".pkl")
- ]
- if self.build_graph_cfg is not None:
+ self.molecules = [
+ osp.join(molecule_cache_path, f"{i:010d}.pkl")
+ for i in range(self.num_samples)
+ ]
+ if build_graph_cfg is not None:
self.graphs = [
- osp.join(self.graphs_dir, f)
- for f in sorted(os.listdir(self.graphs_dir))
- if f.endswith(".pkl")
+ osp.join(graph_cache_path, f"{i:010d}.pkl")
+ for i in range(self.num_samples)
]
else:
self.graphs = None
- logger.info(f"Loading properties {self.property_names} into memory...")
-
- self.property_data = {}
- for pname in self.property_names:
-
- # Determine the actual file name: use the mapping table if available; otherwise, use the configuration name itself
- file_name = PROPERTY_FILE_MAP.get(pname, pname)
-
- file_path = osp.join(self.props_dir, f"{file_name}.pkl")
-
- if not osp.exists(file_path):
- raise FileNotFoundError(
- f"[QM9 Map Error]can't find the file: {file_path}. "
- f"(require label: {pname}, actually find the file: {file_name}.pkl)"
- )
-
- # Although the name of file is lumo.pkl,The key stored in the dictionary is still pname (such as energy_per_atom)
- self.property_data[pname] = self._load_pickle(file_path)
-
- # self.property_data = {
- # pname: self._load_pickle(osp.join(self.props_dir, f"{pname}.pkl"))
- # for pname in self.property_names
- # }
-
-
- # Sort files to ensure consistency across distributed ranks
- # 5) Filter invalid data based on properties and graphs
- if self.filter_unvalid:
- self._filter_by_properties()
- if self.graphs is not None:
- self._filter_by_graphs()
- # 6) Ensure data length consistency across all arrays
- self._ensure_length_consistency()
-
- self.num_samples = len(self.structures)
- logger.info(f"Final QM9Dataset samples: {self.num_samples}")
-
-
- def _prepare_structures_and_properties(self, raw_file_path: str):
- """
- Check if structures and properties are cached; rebuild if missing
- or overwrite is True.
- """
-
- num_cached = self._count_files(self.structures_dir)
-
- # Check if all property files exist
- props_exist = all(
- osp.exists(osp.join(self.props_dir, f"{p}.pkl"))
- for p in self.property_names
- )
-
- # Use a completion flag to ensure the previous build was successful
- struct_done_flag = osp.join(self.structures_dir, "completed.flag")
- is_complete = osp.exists(struct_done_flag)
-
-
- should_build = (
- self.overwrite or num_cached == 0 or not props_exist or not is_complete
- )
-
- if should_build:
- if dist.get_rank() == 0:
- logger.info("Building structures and properties from raw QM9 file...")
-
- # Clean old data to prevent mixing files
- self._clean_dir(self.structures_dir)
- self._clean_dir(self.props_dir)
-
- self._build_structures_and_properties(
- raw_file_path, self.structures_dir, self.props_dir
- )
-
- # Write completion flag
- with open(struct_done_flag, "w") as f:
- f.write("done")
- else:
- logger.info(f"Using cached structures ({num_cached}) and properties.")
-
-
- def _prepare_graphs(self):
- """
- Check if graphs are cached; rebuild if missing, incomplete,
- or overwrite is True.
- """
- num_structs = self._count_files(self.structures_dir)
- num_graphs = self._count_files(self.graphs_dir)
-
- # Use a completion flag for graphs
- graph_done_flag = osp.join(self.graphs_dir, "completed.flag")
- is_complete = osp.exists(graph_done_flag)
-
- # Condition: Not overwrite, marked complete, and counts match
- if not self.overwrite and is_complete and num_graphs == num_structs:
- logger.info(f"Using cached graphs ({num_graphs}).")
- return
-
- logger.info(
- f"Rebuilding graphs. (Structs: {num_structs}, Graphs: {num_graphs}, "
- f"Complete: {is_complete}, Overwrite: {self.overwrite})"
- )
-
- self._clean_dir(self.graphs_dir)
- converter = build_graph_converter(self.build_graph_cfg)
- self._build_graphs(converter, self.structures_dir, self.graphs_dir)
-
- # Write completion flag
- with open(graph_done_flag, "w") as f:
- f.write("done")
-
- def _build_graphs(self, converter, structures_dir: str, graphs_dir: str) -> None:
- """
- Builds graph objects from structures using a SINGLE global progress bar.
-
- This method processes structures in batches to manage memory usage, while
- providing a unified progress visualization.
- """
- import gc
- import sys
-
- # Context manager to temporarily suppress stderr
- class SuppressStderr:
- def __init__(self):
- self.null_fds = [os.open(os.devnull, os.O_RDWR)]
- self.save_fds = [os.dup(2)] # Backup stderr (fd 2)
-
- def __enter__(self):
- # Redirect stderr to devnull
- os.dup2(self.null_fds[0], 2)
-
- def __exit__(self, *_):
- # Restore stderr
- os.dup2(self.save_fds[0], 2)
- for fd in self.null_fds + self.save_fds:
- os.close(fd)
-
- # Get file list
- files = sorted([f for f in os.listdir(structures_dir) if f.endswith(".pkl")])
- total = len(files)
- if total == 0:
- logger.warning("No structures found to convert!")
- return
-
- batch_size = 2000 # Define batch size
-
- logger.info(f"Converting {total} structures to graphs...")
-
- # 1. Create global progress bar
- pbar = tqdm(total=total, desc="Graph Conversion", unit="sample")
-
- for start_idx in range(0, total, batch_size):
- end_idx = min(start_idx + batch_size, total)
- batch_files = files[start_idx:end_idx]
-
- try:
- # Load structures for current batch
- structures = [
- self._load_pickle(osp.join(structures_dir, f)) for f in batch_files
+ assert (
+ len(self.molecules) == self.num_samples
+ ), "The number of molecules must be equal to the number of samples."
+ assert (
+ self.graphs is None or len(self.graphs) == self.num_samples
+ ), "The number of graphs must be equal to the number of samples."
+
+ if filter_unvalid:
+ self.filter_unvalid_by_property()
+
+ def read_data(self, path: str):
+ """Read the data from the given csv path."""
+ data = pd.read_csv(path)
+ logger.info(f"Read {len(data)} molecules from {path}")
+
+ data = {key: data[key].tolist() for key in data if "Unnamed" not in key}
+ data["molecule"] = data.pop("standard_xyz") # adapted for this qm9 split file
+ data["id"] = data["molecule_id"] # adapted for this qm9 split file
+ num_samples = 0
+ for key in data:
+ num_samples = max(num_samples, len(data[key]))
+ return data, num_samples
+
+ def filter_unvalid_by_property(self):
+ for property_name in self.property_names:
+ data = self.property_data[property_name]
+ reserve_idx = []
+ for i, data_item in enumerate(data):
+ if isinstance(data_item, str) or (
+ data_item is not None and not math.isnan(data_item)
+ ):
+ reserve_idx.append(i)
+ for key in self.property_data.keys():
+ self.property_data[key] = [
+ self.property_data[key][i] for i in reserve_idx
]
+ for key in self.raw_data.keys():
+ self.raw_data[key] = [self.raw_data[key][i] for i in reserve_idx]
+ self.molecules = [self.molecules[i] for i in reserve_idx]
+ if self.graphs is not None:
+ self.graphs = [self.graphs[i] for i in reserve_idx]
+ logger.warning(
+ f"Filter out {len(reserve_idx)} samples with valid properties: "
+ f"{property_name}"
+ )
+ self.num_samples = len(self.molecules)
+ logger.warning(f"Remaining {self.num_samples} samples after filtering.")
- # 2. Convert to graphs (Suppress internal progress bars if any)
- try:
- with SuppressStderr():
- graphs = converter(structures)
- except Exception:
- # Fallback if low-level FD manipulation fails
- graphs = converter(structures)
-
- # Save graphs
- for f, g in zip(batch_files, graphs):
- self._save_pickle(osp.join(graphs_dir, f), g)
+ def read_property_data(self, data: Dict, property_names: list[str]):
+ property_data = {
+ property_name: data[property_name] for property_name in property_names
+ }
+ return property_data
- # 3. Update global progress bar
- pbar.update(len(batch_files))
+ def save_to_cache(self, cache_path: str, data: Any):
+ with open(cache_path, "wb") as f:
+ pickle.dump(data, f)
- except Exception as e:
- # Restore stderr to print error
- sys.stderr = sys.__stderr__
- logger.warning(f"Batch {start_idx}-{end_idx} failed: {e}")
-
- finally:
- if "structures" in locals():
- del structures
- if "graphs" in locals():
- del graphs
- gc.collect()
-
- pbar.close()
- logger.info("Graph conversion completed.")
-
- def _ensure_length_consistency(self):
- """
- Ensures consistency in length across structures, graphs, and all
- property arrays. Truncates data to the minimum length found.
- """
- lengths = [len(self.structures)]
- if self.graphs is not None:
- lengths.append(len(self.graphs))
- for p in self.property_names:
- lengths.append(len(self.property_data[p]))
-
- min_len = min(lengths)
-
- if any(length != min_len for length in lengths):
- logger.warning(
- f"Data length mismatch detected (lengths={lengths}). "
- f"Truncating to minimum length: {min_len}."
- )
- self.structures = self.structures[:min_len]
- if self.graphs is not None:
- self.graphs = self.graphs[:min_len]
- for p in self.property_names:
- self.property_data[p] = self.property_data[p][:min_len]
+ def load_from_cache(self, cache_path: str):
+ if osp.exists(cache_path):
+ with open(cache_path, "rb") as f:
+ data = pickle.load(f)
+ return data
+ raise FileNotFoundError(f"No such file or directory: {cache_path}")
+ def __getitem__(self, idx: int):
+ data = {}
- def _clean_dir(self, directory: str):
- """Cleans a directory by removing all .pkl and .flag files."""
- for f in os.listdir(directory):
- if f.endswith(".pkl") or f.endswith(".flag"):
- try:
- os.remove(osp.join(directory, f))
- except OSError:
- pass
-
-
- def _ensure_raw_data(self) -> str:
- """
- downloading self.url -> self.raw_xyz_path
- """
- # 1. if the final file exists , return.
- if osp.exists(self.raw_xyz_path):
- return self.raw_xyz_path
-
- # 2. prepare the path to download
- tar_filename = "qm9_raw.tar.bz2"
- tar_path = osp.join(self.raw_dir, tar_filename)
-
- # 3. downloading logic
- if not osp.exists(tar_path):
- if dist.get_rank() == 0:
- logger.info(f"Downloading QM9 from {self.url}...")
- import urllib.request
- try:
- urllib.request.urlretrieve(self.url, tar_path)
- except Exception as e:
- raise RuntimeError(f"Download failed: {e}")
- if dist.is_initialized():
- dist.barrier()
-
- # 4. extacting logic
- if dist.get_rank() == 0:
- logger.info("Extracting QM9...")
- import tarfile
- try:
- with tarfile.open(tar_path, "r:bz2") as tar:
- tar.extractall(path=self.raw_dir)
- except Exception as e:
- raise RuntimeError(f"Extraction failed: {e}")
-
- if dist.is_initialized():
- dist.barrier()
-
- # 5. final check
- # Case A：single file exists.
- if osp.exists(self.raw_xyz_path):
- return self.raw_xyz_path
-
-
- # Case B：merge these .xyz files into a big file.
- xyz_files = [f for f in os.listdir(self.raw_dir) if f.endswith(".xyz") and f != "dsgdb9nsd.xyz"]
- if len(xyz_files) > 0:
- logger.info(f"Found {len(xyz_files)} xyz files, merging into dsgdb9nsd.xyz...")
- merged_path = self.raw_xyz_path
-
- if osp.exists(merged_path):
- os.remove(merged_path)
-
- with open(merged_path, "w") as fout: # use "w" model to rewrite
- for fname in tqdm(sorted(xyz_files), desc="Merging XYZ files"):
- full_path = osp.join(self.raw_dir, fname)
- try:
- with open(full_path, "r") as fin:
- lines = fin.readlines()
-
- if not lines: continue
- natoms = int(lines[0].strip())
-
- # 1. Number of atoms written
- fout.write(f"{natoms}\n")
- # 2. Write attribute line
- prop_line = lines[1].replace('*^', 'e').replace('\t', ' ')
- fout.write(prop_line)
- # 3. Write coordinate lines (only take natoms lines)
- for i in range(2, 2 + natoms):
- coord_line = lines[i].replace('*^', 'e').replace('\t', ' ')
- fout.write(coord_line)
-
- except Exception as e:
- logger.warning(f"Error processing {fname}: {e}")
- continue
- return merged_path
- # Case C: None
- raise RuntimeError(
- f"Decompression is complete, but I couldn't find dsgdb9nsd.xyz or any .xyz files under {self.raw_dir}!"
- "Please check what files are actually included in the downloaded compressed package."
- )
-
- def _count_files(self, directory: str) -> int:
- """Counts the number of .pkl files in a directory."""
- try:
- return len([n for n in os.listdir(directory) if n.endswith(".pkl")])
- except Exception:
- return 0
-
- @staticmethod
- def _save_pickle(path: str, obj: Any) -> None:
- with open(path, "wb") as f:
- pickle.dump(obj, f)
-
- @staticmethod
- def _load_pickle(path: str) -> Any:
- with open(path, "rb") as f:
- return pickle.load(f)
-
- def _build_structures_and_properties(self, raw_path: str, struct_dir: str, prop_dir: str) -> None:
- """
- Core Constructor Function:analyse XYZ -> Pymatgen Structure -> Pickle
- """
- logger.info(f"Parsing {raw_path} using ASE...")
-
- # 1. Read all data into memory (QM9 is about 100MB, which can easily fit into memory)
- atoms_collection = read(raw_path, index=':')
-
- # 2. Read text lines to parse attributes (ASE attribute parsing is sometimes unreliable; manual parsing is more stable)
- with open(raw_path, 'r') as f:
- lines = f.readlines()
-
- prop_buffers = defaultdict(list)
- current_line = 0
- valid_count = 0
-
- total = len(atoms_collection)
- pbar = tqdm(total=total, desc="Processing QM9")
-
- for i, atoms in enumerate(atoms_collection):
- try:
- num_atoms = len(atoms)
- prop_line = lines[current_line + 1]
-
- # clean the property line.
- prop_line_cleaned = prop_line.replace('*^', 'e').replace('\t', ' ')
- raw_vals = prop_line_cleaned.split()
-
- vals_float = []
- for val_str in raw_vals:
- try:
- vals_float.append(float(val_str))
- except ValueError:
- # for strings like 'gdb'
- vals_float.append(0.0)
-
- # Mapping Attribute
- for k, key in enumerate(self.PROP_ORDER):
- if key in ['tag', 'index']:
- continue
-
- # Alignment index: k=2 is 'A', corresponding to raw_vals[2]
- if k < len(vals_float):
- prop_buffers[key].append(vals_float[k])
- else:
- prop_buffers[key].append(np.nan)
-
- # --- B. building Structure (Fake crystal cell) ---
- # Set up a large box to prevent the model from reporting errors due to the absence of cells
- atoms.set_cell([20.0, 20.0, 20.0])
- atoms.center()
- atoms.pbc = True
- structure = AseAtomsAdaptor.get_structure(atoms)
-
- # --- C. save the struncture ---
- self._save_pickle(osp.join(struct_dir, f"{i:06d}.pkl"), structure)
-
- # Update pointer
- current_line += (num_atoms + 2)
- valid_count += 1
- pbar.update(1)
-
- except Exception as e:
- logger.warning(f"Error processing molecule {i}: {e}. Skipping block.")
- current_line += (len(atoms) + 2)
- continue
-
- pbar.close()
-
- if valid_count == 0:
- raise RuntimeError("No valid samples processed from QM9 file!")
-
- # --- D.Save attribute array ---
- logger.info("Saving property arrays...")
- for key, val_list in prop_buffers.items():
- self._save_pickle(
- osp.join(prop_dir, f"{key}.pkl"),
- np.array(val_list, dtype=np.float32)
+ if self.graphs is None:
+ raise ValueError(
+ "QM9Dataset requires build_graph_cfg to return model-ready samples."
)
-
- def _filter_by_properties(self) -> None:
- """
- Filter out samples that contain invalid property values (e.g., NaN, Inf).
- Operation is performed in-memory.
- """
- if not self.property_names:
- return
-
- total = len(self.structures)
- keep = []
-
- # Check properties for each sample
- for i in range(total):
- is_valid = True
- for pname in self.property_names:
- val = self.property_data[pname][i]
- if val is None:
- is_valid = False
- break
-
- # Check for NaN/Inf in scalars
- if isinstance(val, (float, int, np.floating, np.integer)):
- if np.isnan(val) or np.isinf(val):
- is_valid = False
- break
- # Check for NaN/Inf in arrays/lists
- elif isinstance(val, (list, np.ndarray)):
- arr = np.asarray(val)
- if not np.all(np.isfinite(arr)):
- is_valid = False
- break
-
- if is_valid:
- keep.append(i)
-
- if len(keep) < total:
- logger.warning(
- f"Filtering: Dropping {total - len(keep)} samples "
- "due to invalid properties."
- )
- self.structures = [self.structures[i] for i in keep]
- if self.graphs:
- self.graphs = [self.graphs[i] for i in keep]
- # stay property_data 为 numpy.ndarray（Indexing with numpy）
- for pname in self.property_names:
- arr = self.property_data[pname]
- # arr may already be numpy array; this keeps dtype and shape consistent
- self.property_data[pname] = arr[keep]
-
- def _filter_by_graphs(self) -> None:
- """
- Filter out samples with invalid or missing graphs.
- Since graphs are rebuilt fully if mismatch occurs, this is mostly a
- sanity check.
- """
- pass
-
- def __len__(self) -> int:
- return self.num_samples
+ graph = self.graphs[idx]
+ if isinstance(graph, str):
+ graph = self.load_from_cache(graph)
+ data["graph"] = graph
- def __getitem__(self, idx: int) -> Dict[str, Any]:
- data = {}
- # 1. loading the info
- if self.graphs is not None:
- data["graph"] = self._load_pickle(self.graphs[idx])
- else:
- struct = self._load_pickle(self.structures[idx])
- # turn into dictionary format
- data["pos"] = np.array(struct.cart_coords, dtype='float32')
- data["atomic_numbers"] = np.array([s.specie.Z for s in struct], dtype='int64')
- data["cell"] = np.array(struct.lattice.matrix, dtype='float32')
- data["natoms"] = len(struct)
- data["pbc"] = np.array([True, True, True], dtype=bool)
-
- # 2. loading the properties
- for pname in self.property_names:
- val = self.property_data[pname][idx]
- # data[pname] = np.array([val], dtype='float32')
- if pname == 'lumo':
- data['energy_per_atom'] = np.array([val], dtype='float32')
+ for property_name in self.property_names:
+ if property_name in self.property_data:
+ data[property_name] = np.array(
+ [self.property_data[property_name][idx]], dtype=np.float32
+ )
else:
- data[pname] = np.array([val], dtype='float32')
-
- # 3. data transforms
- if self.transforms is not None:
- data = self.transforms(data)
-
- return data
\ No newline at end of file
+ raise KeyError(f"Property {property_name} not found.")
+ data["id"] = self.raw_data["id"][idx]
+ data = self.transforms(data) if self.transforms is not None else data
+
+ return data
+
+ def __len__(self):
+ return self.num_samples
diff --git a/ppmat/datasets/script_split_qm9.py b/ppmat/datasets/script_split_qm9.py
new file mode 100644
index 00000000..b3fa6ab3
--- /dev/null
+++ b/ppmat/datasets/script_split_qm9.py
@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+"""Export QM9 into train.csv, val.csv, and test.csv.
+
+Output rows are one molecule each. The split policy matches the earlier export:
+remove Figshare's 3,054 uncharacterized molecule IDs, then shuffle the remaining
+130,831 molecule IDs with seed 42 and split as 110000/10000/10831.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import random
+import re
+import urllib.request
+from pathlib import Path
+
+
+UNCHARACTERIZED_URL = "https://ndownloader.figshare.com/files/3195404"
+
+PROPERTY_COLUMNS = [
+ "A",
+ "B",
+ "C",
+ "mu",
+ "alpha",
+ "homo",
+ "lumo",
+ "gap",
+ "r2",
+ "zpve",
+ "U0",
+ "U",
+ "H",
+ "G",
+ "Cv",
+]
+
+HEADER = [
+ "file_name",
+ "raw_file_content",
+ "standard_xyz",
+ "mulliken_xyz",
+ "molecule_id",
+ "num_atoms",
+ *PROPERTY_COLUMNS,
+ "vibrational_frequencies",
+ "canonical_smiles",
+ "isomeric_smiles",
+ "canonical_inchi",
+ "isomeric_inchi",
+]
+
+
+def natural_key(path: Path) -> int:
+ match = re.search(r"(\d+)$", path.stem)
+ if not match:
+ raise ValueError(f"Cannot infer molecule id from {path.name}")
+ return int(match.group(1))
+
+
+def parse_float_text(value: str) -> str:
+ """Normalize rare Mathematica-style exponents while keeping CSV text stable."""
+ return str(float(value.replace("*^", "e")))
+
+
+def parse_xyz_coord_text(value: str) -> str:
+ """Format XYZ coordinates as plain decimals for RDKit compatibility."""
+ return f"{float(value.replace('*^', 'e')):.10f}"
+
+
+def parse_charge_text(value: str) -> str:
+ """Format Mulliken partial charges as plain decimals."""
+ return f"{float(value.replace('*^', 'e')):.10f}"
+
+
+def ensure_uncharacterized(path: Path) -> Path:
+ if not path.exists() or path.stat().st_size == 0:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ urllib.request.urlretrieve(UNCHARACTERIZED_URL, path)
+ return path
+
+
+def load_bad_ids(path: Path) -> set[int]:
+ bad = set()
+ for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+ parts = line.split()
+ if parts and parts[0].isdigit():
+ bad.add(int(parts[0]))
+ if len(bad) != 3054:
+ raise ValueError(f"Expected 3054 uncharacterized ids, got {len(bad)}")
+ return bad
+
+
+def split_ids(ids: list[int], seed: int) -> dict[int, str]:
+ shuffled = ids[:]
+ random.Random(seed).shuffle(shuffled)
+ train = set(shuffled[:110_000])
+ val = set(shuffled[110_000:120_000])
+ return {
+ molecule_id: "train"
+ if molecule_id in train
+ else "val"
+ if molecule_id in val
+ else "test"
+ for molecule_id in ids
+ }
+
+
+def parse_qm9_file(path: Path) -> dict[str, object]:
+ raw = path.read_text(encoding="utf-8", errors="replace")
+ lines = raw.splitlines()
+ num_atoms = int(lines[0].strip())
+
+ props = lines[1].split()
+ if len(props) < 17:
+ raise ValueError(f"Malformed property line in {path.name}")
+ molecule_id = int(props[1])
+ if molecule_id != natural_key(path):
+ raise ValueError(f"{path.name} contains molecule id {molecule_id}")
+
+ property_values = dict(zip(PROPERTY_COLUMNS, [parse_float_text(x) for x in props[2:17]]))
+ comment_line = f"gdb {molecule_id}"
+
+ standard_atom_lines = []
+ mulliken_atom_lines = []
+ for line in lines[2 : 2 + num_atoms]:
+ parts = line.split()
+ if len(parts) != 5:
+ raise ValueError(f"Malformed atom line in {path.name}: {line!r}")
+ element = parts[0]
+ x, y, z = [parse_xyz_coord_text(x) for x in parts[1:4]]
+ charge = parse_charge_text(parts[4])
+ standard_atom_lines.append(f"{element}\t{x}\t{y}\t{z}")
+ mulliken_atom_lines.append(f"{element}\t{x}\t{y}\t{z}\t{charge}")
+
+ standard_xyz = "\n".join([str(num_atoms), comment_line, *standard_atom_lines])
+ mulliken_xyz = "\n".join([str(num_atoms), comment_line, *mulliken_atom_lines])
+
+ freq_idx = 2 + num_atoms
+ frequencies = " ".join(parse_float_text(x) for x in lines[freq_idx].split())
+ smiles = lines[freq_idx + 1].split()
+ inchi = lines[freq_idx + 2].split()
+
+ return {
+ "file_name": path.name,
+ "raw_file_content": raw,
+ "standard_xyz": standard_xyz,
+ "mulliken_xyz": mulliken_xyz,
+ "molecule_id": molecule_id,
+ "num_atoms": num_atoms,
+ **property_values,
+ "vibrational_frequencies": frequencies,
+ "canonical_smiles": smiles[0] if len(smiles) > 0 else "",
+ "isomeric_smiles": smiles[1] if len(smiles) > 1 else "",
+ "canonical_inchi": inchi[0] if len(inchi) > 0 else "",
+ "isomeric_inchi": inchi[1] if len(inchi) > 1 else "",
+ }
+
+
+def export(input_dir: Path, output_dir: Path, uncharacterized: Path, seed: int) -> dict[str, int]:
+ output_dir.mkdir(parents=True, exist_ok=True)
+ bad_ids = load_bad_ids(ensure_uncharacterized(uncharacterized))
+ files = sorted(input_dir.glob("*.xyz"), key=natural_key)
+ ids = [natural_key(path) for path in files if natural_key(path) not in bad_ids]
+ split_lookup = split_ids(ids, seed)
+
+ handles = {}
+ writers = {}
+ counts = {"train": 0, "val": 0, "test": 0}
+ try:
+ for split in ["train", "val", "test"]:
+ handle = (output_dir / f"{split}.csv").open("w", newline="", encoding="utf-8")
+ writer = csv.DictWriter(handle, fieldnames=HEADER)
+ writer.writeheader()
+ handles[split] = handle
+ writers[split] = writer
+
+ for path in files:
+ molecule_id = natural_key(path)
+ split = split_lookup.get(molecule_id)
+ if split is None:
+ continue
+ writers[split].writerow(parse_qm9_file(path))
+ counts[split] += 1
+ finally:
+ for handle in handles.values():
+ handle.close()
+ return counts
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--input-dir", type=Path, default=Path("./qm9.xyz"))
+ parser.add_argument("--output-dir", type=Path, default=Path("./qm9_split"))
+ parser.add_argument(
+ "--uncharacterized",
+ type=Path,
+ default=Path("./qm9_split_scripts/uncharacterized.txt"),
+ )
+ parser.add_argument("--seed", type=int, default=42)
+ args = parser.parse_args()
+ counts = export(args.input_dir, args.output_dir, args.uncharacterized, args.seed)
+ print(counts)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/ppmat/datasets/split_gmtnet_dataset.py b/ppmat/datasets/split_gmtnet_dataset.py
new file mode 100644
index 00000000..242c75ad
--- /dev/null
+++ b/ppmat/datasets/split_gmtnet_dataset.py
@@ -0,0 +1,243 @@
+"""Generate or verify the fixed-seed GMTNet dielectric split.
+
+This offline utility only manages the canonical index split. It neither
+downloads nor converts datasets, and it never overwrites a split unless
+``--force`` is supplied explicitly.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any
+
+
+EXPECTED_SOURCE_SHA256 = (
+ "5a2198f51f4a7f9aa26fa6be60ed65db0ecc0a13646d63399e897b8168dcbb4d"
+)
+EXPECTED_NORMALIZED_SHA256 = (
+ "eb0b9516c937575afe3f20a0f88953724abfcde6de07c4af15248468598b349f"
+)
+EXPECTED_RECORD_COUNT = 4713
+EXPECTED_SPLIT_SEED = 32
+EXPECTED_SPLIT_SIZES = {"train": 3770, "val": 471, "test": 472}
+_SPLIT_NAMES = tuple(EXPECTED_SPLIT_SIZES)
+
+
+def _split_indices_sha256(indices: list[int]) -> str:
+ """Return the canonical hash for one ordered index list."""
+ encoded = json.dumps(indices, separators=(",", ":"), ensure_ascii=True).encode(
+ "utf-8"
+ )
+ return hashlib.sha256(encoded).hexdigest()
+
+
+def _validate_indices(indices: Any, split_name: str) -> list[int]:
+ """Validate one fixed-size partition without changing its order."""
+ expected_size = EXPECTED_SPLIT_SIZES[split_name]
+ if not isinstance(indices, list) or len(indices) != expected_size:
+ raise ValueError(f"{split_name}_indices length must be {expected_size}.")
+ if any(type(index) is not int for index in indices):
+ raise ValueError(f"{split_name}_indices must contain JSON integers only.")
+ if any(index < 0 or index >= EXPECTED_RECORD_COUNT for index in indices):
+ raise ValueError(f"{split_name}_indices contains an out-of-range index.")
+ if len(indices) != len(set(indices)):
+ raise ValueError(f"{split_name}_indices contains duplicate indices.")
+ return indices
+
+
+def validate_split_data(split_data: Any) -> dict[str, list[int]]:
+ """Validate the canonical GMTNet split schema and ordered partitions."""
+ required_fields = {
+ "schema_version",
+ "source_original_dataset_sha256",
+ "normalized_schema_version",
+ "normalized_dataset_sha256",
+ "num_records",
+ "seed",
+ "generation_method",
+ "torch_version",
+ "split_sizes",
+ "split_indices_sha256",
+ "train_indices",
+ "val_indices",
+ "test_indices",
+ }
+ if not isinstance(split_data, dict) or set(split_data) != required_fields:
+ raise ValueError("Split JSON fields do not match the canonical schema.")
+ expected_values = {
+ "schema_version": 1,
+ "source_original_dataset_sha256": EXPECTED_SOURCE_SHA256,
+ "normalized_schema_version": 1,
+ "normalized_dataset_sha256": EXPECTED_NORMALIZED_SHA256,
+ "num_records": EXPECTED_RECORD_COUNT,
+ "seed": EXPECTED_SPLIT_SEED,
+ "generation_method": "torch.utils.data.random_split",
+ "split_sizes": EXPECTED_SPLIT_SIZES,
+ }
+ for field_name, expected_value in expected_values.items():
+ if split_data[field_name] != expected_value:
+ raise ValueError(f"Split JSON {field_name} is invalid.")
+ if not isinstance(split_data["torch_version"], str) or not split_data["torch_version"]:
+ raise ValueError("Split JSON torch_version is invalid.")
+
+ split_indices = {
+ split_name: _validate_indices(
+ split_data[f"{split_name}_indices"], split_name
+ )
+ for split_name in _SPLIT_NAMES
+ }
+ all_indices = [index for indices in split_indices.values() for index in indices]
+ if len(all_indices) != EXPECTED_RECORD_COUNT or len(set(all_indices)) != len(
+ all_indices
+ ):
+ raise ValueError("Split partitions must be disjoint and cover every record.")
+ if set(all_indices) != set(range(EXPECTED_RECORD_COUNT)):
+ raise ValueError("Split partitions do not cover the complete dataset.")
+ if split_indices["test"][:3] != [747, 1423, 1322]:
+ raise ValueError("Split JSON test_indices fixed prefix is invalid.")
+
+ hashes = split_data["split_indices_sha256"]
+ if not isinstance(hashes, dict) or set(hashes) != set(_SPLIT_NAMES):
+ raise ValueError("Split JSON split_indices_sha256 fields are invalid.")
+ for split_name, indices in split_indices.items():
+ if hashes[split_name] != _split_indices_sha256(indices):
+ raise ValueError(f"Split JSON {split_name}_indices SHA256 is invalid.")
+ return split_indices
+
+
+def load_and_validate_split(split_path: str | Path) -> dict[str, list[int]]:
+ """Load and validate a split JSON file."""
+ path = Path(split_path).expanduser().resolve()
+ if not path.is_file():
+ raise FileNotFoundError(f"split_path must be a regular file: {path}")
+ with path.open("r", encoding="utf-8") as handle:
+ return validate_split_data(json.load(handle))
+
+
+def _require_torch_for_generation():
+ try:
+ import torch
+ except ModuleNotFoundError as error:
+ raise RuntimeError(
+ "Generating the historical random_split order requires PyTorch. "
+ "Use --verify to validate an existing split without PyTorch."
+ ) from error
+ return torch
+
+
+def generate_split(seed: int = EXPECTED_SPLIT_SEED) -> dict[str, Any]:
+ """Reproduce the official PyTorch ``random_split`` index order."""
+ if seed != EXPECTED_SPLIT_SEED:
+ raise ValueError(f"seed must be {EXPECTED_SPLIT_SEED}.")
+ torch = _require_torch_for_generation()
+ generator = torch.Generator()
+ generator.manual_seed(seed)
+ train_subset, val_subset, test_subset = torch.utils.data.random_split(
+ range(EXPECTED_RECORD_COUNT),
+ list(EXPECTED_SPLIT_SIZES.values()),
+ generator=generator,
+ )
+ split_indices = {
+ "train": list(train_subset.indices),
+ "val": list(val_subset.indices),
+ "test": list(test_subset.indices),
+ }
+ split_data = {
+ "schema_version": 1,
+ "source_original_dataset_sha256": EXPECTED_SOURCE_SHA256,
+ "normalized_schema_version": 1,
+ "normalized_dataset_sha256": EXPECTED_NORMALIZED_SHA256,
+ "num_records": EXPECTED_RECORD_COUNT,
+ "seed": seed,
+ "generation_method": "torch.utils.data.random_split",
+ "torch_version": str(torch.__version__),
+ "split_sizes": EXPECTED_SPLIT_SIZES,
+ "split_indices_sha256": {
+ split_name: _split_indices_sha256(indices)
+ for split_name, indices in split_indices.items()
+ },
+ "train_indices": split_indices["train"],
+ "val_indices": split_indices["val"],
+ "test_indices": split_indices["test"],
+ }
+ validate_split_data(split_data)
+ return split_data
+
+
+def _write_split(split_data: dict[str, Any], output_path: Path, force: bool) -> None:
+ output_path = output_path.expanduser().resolve()
+ if output_path.exists() and not force:
+ raise FileExistsError(
+ f"Refusing to overwrite existing split JSON: {output_path}. Use --force."
+ )
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+ descriptor, temporary_name = tempfile.mkstemp(
+ prefix=f".{output_path.name}.", suffix=".tmp", dir=output_path.parent
+ )
+ temporary_path = Path(temporary_name)
+ try:
+ with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+ json.dump(split_data, handle, indent=2, sort_keys=True)
+ handle.write("\n")
+ handle.flush()
+ os.fsync(handle.fileno())
+ os.replace(temporary_path, output_path)
+ except Exception:
+ temporary_path.unlink(missing_ok=True)
+ raise
+
+
+def _build_parser() -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(
+ description="Generate or verify the fixed-seed GMTNet dielectric split."
+ )
+ action = parser.add_mutually_exclusive_group(required=True)
+ action.add_argument("--verify", metavar="SPLIT_JSON", help="Validate a split JSON.")
+ action.add_argument("--output", metavar="SPLIT_JSON", help="Write a generated split JSON.")
+ parser.add_argument("--seed", type=int, default=EXPECTED_SPLIT_SEED)
+ parser.add_argument(
+ "--record-count",
+ type=int,
+ default=EXPECTED_RECORD_COUNT,
+ help=f"Must be {EXPECTED_RECORD_COUNT}; validates the source record count.",
+ )
+ parser.add_argument(
+ "--force", action="store_true", help="Allow overwriting --output."
+ )
+ return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+ """Run the split CLI."""
+ args = _build_parser().parse_args(argv)
+ try:
+ if args.record_count != EXPECTED_RECORD_COUNT:
+ raise ValueError(
+ f"record_count must be {EXPECTED_RECORD_COUNT}; "
+ "the source dataset is incomplete or incompatible."
+ )
+ if args.verify:
+ split_indices = load_and_validate_split(args.verify)
+ print(
+ json.dumps(
+ {"status": "ok", "sizes": {key: len(value) for key, value in split_indices.items()}},
+ sort_keys=True,
+ )
+ )
+ else:
+ _write_split(generate_split(args.seed), Path(args.output), args.force)
+ print(json.dumps({"status": "ok", "output": str(Path(args.output).resolve())}))
+ except Exception as error:
+ print(f"ERROR: {type(error).__name__}: {error}", file=sys.stderr)
+ return 1
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py
index d258bd2a..99f10251 100644
--- a/ppmat/models/__init__.py
+++ b/ppmat/models/__init__.py
@@ -29,24 +29,29 @@
from ppmat.models.common.graph_converter import CrystalNN
from ppmat.models.common.graph_converter import FindPointsInSpheres
from ppmat.models.common.graph_converter import MolecularGraphConverter
+from ppmat.models.common.graph_converter import RadiusGraphConverter
from ppmat.models.diffcsp.diffcsp import DiffCSP
from ppmat.models.diffnmr.diffnmr import DiffNMR
from ppmat.models.diffnmr.diffnmr import DiffPrior
from ppmat.models.diffnmr.diffnmr import MolecularGraphFormer
from ppmat.models.diffnmr.diffnmr import NMRNetCLIP
from ppmat.models.dimenetpp.dimenetpp import DimeNetPlusPlus
+from ppmat.models.gmtnet.gmtnet import GMTNet
+from ppmat.models.gmtnet.gmtnet_graph_converter import GMTNetGraphConverter
+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.sfin.sfin import SFIN
+from ppmat.models.spherenet.spherenet import SphereNet
from ppmat.utils import download
from ppmat.utils import logger
from ppmat.utils import save_load
+
__all__ = [
"iComformer",
"ComformerGraphConverter",
@@ -56,12 +61,15 @@
"MatterGen",
"MatterGenWithCondition",
"DimeNetPlusPlus",
+ "GMTNet",
+ "GMTNetGraphConverter",
"CrystalNN",
"CHGNetGraphConverter",
"CHGNet",
"M3GNetGraphConvertor",
"M3GNet",
"MolecularGraphConverter",
+ "RadiusGraphConverter",
"MolecularGraphFormer",
"NMRNetCLIP",
"DiffPrior",
@@ -69,6 +77,7 @@
"InfGCN",
"MatENO",
"SFIN",
+ "SphereNet",
]
# Warning: The key of the dictionary must be consistent with the file name of the value
@@ -117,6 +126,26 @@
"sfin_haadf_detect": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_haadf_detect.zip",
"sfin_bf_enhance": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_bf_enhance.zip",
"sfin_bf_detect": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/spectrum_enhancement/sfin/sfin_bf_detect.zip",
+ "spherenet_qm9_mu": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_mu.zip",
+ "spherenet_qm9_alpha": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_alpha.zip",
+ "spherenet_qm9_homo": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_homo.zip",
+ "spherenet_qm9_lumo": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_lumo.zip",
+ "spherenet_qm9_gap": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_gap.zip",
+ "spherenet_qm9_r2": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_r2.zip",
+ "spherenet_qm9_zpve": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_zpve.zip",
+ "spherenet_qm9_U0": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_U0.zip",
+ "spherenet_qm9_U": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_U.zip",
+ "spherenet_qm9_H": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_H.zip",
+ "spherenet_qm9_G": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_G.zip",
+ "spherenet_qm9_Cv": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/property_prediction/spherenet/spherenet_qm9_Cv.zip",
+ "spherenet_md17_aspirin": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_aspirin.zip",
+ "spherenet_md17_benzene_old": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_benzene_old.zip",
+ "spherenet_md17_ethanol": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_ethanol.zip",
+ "spherenet_md17_malonaldehyde": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_malonaldehyde.zip",
+ "spherenet_md17_naphthalene": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_naphthalene.zip",
+ "spherenet_md17_salicylic": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_salicylic.zip",
+ "spherenet_md17_toluene": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_toluene.zip",
+ "spherenet_md17_uracil": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/spherenet/spherenet_md17_uracil.zip",
}
diff --git a/ppmat/models/common/graph_converter.py b/ppmat/models/common/graph_converter.py
index 2cb13ac7..b71bbc74 100644
--- a/ppmat/models/common/graph_converter.py
+++ b/ppmat/models/common/graph_converter.py
@@ -572,6 +572,255 @@ def __call__(
)
+class RadiusGraphConverter:
+ """Convert RDKit molecules into PGL radius graphs.
+
+ The converter is used by SphereNet-style molecular datasets. It accepts an
+ RDKit ``Mol`` or a list of RDKit ``Mol`` objects and returns PGL graph data
+ with atom numbers, positions, radius edges, and optional SphereNet triplet
+ indices cached in ``graph.edge_feat``.
+
+ Args:
+ cutoff: Neighbor cutoff distance in Angstrom.
+ atom_vocab: Mapping from atomic number to feature index for optional PGL
+ node features.
+ add_self_loops: Whether to add self-loop edges.
+ edge_mode: ``"directed"``, ``"undirected"``, or ``"bidirectional"``.
+ include_distance: Whether to include distance in PGL edge features.
+ include_direction: Whether to include unit direction in PGL edge features.
+ return_triplet_indices: Whether to cache SphereNet triplet indices.
+ num_cpus: Number of CPUs for parallel graph construction.
+ """
+
+ def __init__(
+ self,
+ cutoff: float = 5.0,
+ atom_vocab: Optional[Dict[int, int]] = None,
+ add_self_loops: bool = False,
+ edge_mode: str = "bidirectional",
+ include_distance: bool = True,
+ include_direction: bool = False,
+ return_triplet_indices: bool = False,
+ num_cpus: Optional[int] = None,
+ ) -> None:
+ if atom_vocab is None:
+ atom_vocab = {1: 0, 6: 1, 7: 2, 8: 3, 9: 4}
+ if edge_mode not in {"directed", "undirected", "bidirectional"}:
+ raise ValueError(f"Unknown edge_mode: {edge_mode}")
+
+ self.cutoff = float(cutoff)
+ self.atom_vocab = dict(atom_vocab)
+ self.add_self_loops = add_self_loops
+ self.edge_mode = edge_mode
+ self.include_distance = include_distance
+ self.include_direction = include_direction
+ self.return_triplet_indices = return_triplet_indices
+ self.num_cpus = 1 if num_cpus is None else int(num_cpus)
+
+ def __call__(
+ self, molecule: Union[Chem.Mol, List[Chem.Mol]]
+ ) -> Union[Optional[pgl.Graph], List[Optional[pgl.Graph]]]:
+ if isinstance(molecule, Chem.Mol):
+ graph = self.get_graph_by_radius(molecule)
+ elif isinstance(molecule, list):
+ if self.num_cpus == 1:
+ graph = [self.get_graph_by_radius(mol) for mol in molecule]
+ else:
+ graph = p_map(
+ self.get_graph_by_radius,
+ molecule,
+ num_cpus=self.num_cpus,
+ desc="Building graphs",
+ dynamic_ncols=True,
+ mininterval=0.2,
+ )
+ else:
+ raise TypeError("The input must be an RDKit Mol or a list of them.")
+ return graph
+
+ def get_graph_by_radius(self, molecule: Chem.Mol) -> Optional[pgl.Graph]:
+ if molecule is None:
+ return None
+
+ atomic_numbers, positions = self.get_molecule_array(molecule)
+ num_nodes = positions.shape[0]
+ if num_nodes == 0:
+ return None
+
+ edge_index, distances, directions = self.get_radius_edges(positions)
+ triplet_indices = None
+ if self.return_triplet_indices:
+ triplet_indices = self.get_triplet_indices(edge_index, num_nodes)
+ return self.build_pgl_graph(
+ atomic_numbers,
+ positions,
+ edge_index,
+ distances,
+ directions,
+ triplet_indices,
+ )
+
+ def get_molecule_array(self, molecule: Chem.Mol) -> Tuple[np.ndarray, np.ndarray]:
+ if molecule.GetNumConformers() == 0:
+ raise ValueError("RDKit Mol has no conformer. Cannot build radius graph.")
+
+ atomic_numbers = np.asarray(
+ [atom.GetAtomicNum() for atom in molecule.GetAtoms()], dtype=np.int64
+ )
+ conf = molecule.GetConformer()
+ positions = np.asarray(
+ [
+ [
+ conf.GetAtomPosition(i).x,
+ conf.GetAtomPosition(i).y,
+ conf.GetAtomPosition(i).z,
+ ]
+ for i in range(molecule.GetNumAtoms())
+ ],
+ dtype=np.float32,
+ )
+ return atomic_numbers, positions
+
+ def get_radius_edges(
+ self, positions: np.ndarray
+ ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+ positions = np.asarray(positions, dtype=np.float32)
+ num_nodes = positions.shape[0]
+ if num_nodes == 0:
+ return (
+ np.empty((2, 0), dtype=np.int64),
+ np.empty((0, 1), dtype=np.float32),
+ np.empty((0, 3), dtype=np.float32),
+ )
+
+ diff = positions[None, :, :] - positions[:, None, :]
+ dist = np.linalg.norm(diff, axis=-1)
+ mask = dist < self.cutoff
+ np.fill_diagonal(mask, False)
+ rows, cols = np.where(mask)
+
+ if self.edge_mode == "undirected":
+ keep = rows < cols
+ rows, cols = rows[keep], cols[keep]
+
+ if self.add_self_loops:
+ self_nodes = np.arange(num_nodes, dtype=np.int64)
+ rows = np.concatenate([rows, self_nodes], axis=0)
+ cols = np.concatenate([cols, self_nodes], axis=0)
+
+ if rows.size == 0:
+ return (
+ np.empty((2, 0), dtype=np.int64),
+ np.empty((0, 1), dtype=np.float32),
+ np.empty((0, 3), dtype=np.float32),
+ )
+
+ edge_index = np.stack([rows, cols], axis=0).astype(np.int64)
+ distances = dist[rows, cols].reshape(-1, 1).astype(np.float32)
+ directions = diff[rows, cols].astype(np.float32)
+ directions = directions / np.maximum(distances, 1e-8)
+
+ order = np.argsort(
+ edge_index[0] * max(1, num_nodes) + edge_index[1], kind="mergesort"
+ )
+ edge_index = edge_index[:, order]
+ distances = distances[order]
+ directions = directions[order]
+ return edge_index, distances, directions
+
+ def get_node_feat(
+ self, atomic_numbers: np.ndarray, positions: np.ndarray
+ ) -> Dict[str, np.ndarray]:
+ atomic_numbers = np.asarray(atomic_numbers, dtype=np.int64)
+ node_feat = {
+ "pos": positions.astype(np.float32),
+ "atomic_number": atomic_numbers.reshape(-1, 1),
+ }
+
+ idxs = []
+ for z in atomic_numbers:
+ if int(z) not in self.atom_vocab:
+ return node_feat
+ idxs.append(self.atom_vocab[int(z)])
+ idxs = np.asarray(idxs, dtype=np.int64)
+ node_feat["feat"] = np.eye(len(self.atom_vocab), dtype=np.float32)[idxs]
+ return node_feat
+
+ def build_pgl_graph(
+ self,
+ atomic_numbers: np.ndarray,
+ positions: np.ndarray,
+ edge_index: np.ndarray,
+ distances: np.ndarray,
+ directions: np.ndarray,
+ triplet_indices: Optional[Dict[str, np.ndarray]] = None,
+ ) -> pgl.Graph:
+ edge_feat = {}
+ edge_feats = []
+ if self.include_distance:
+ edge_feat["distance"] = distances
+ edge_feats.append(distances)
+ if self.include_direction:
+ edge_feat["direction"] = directions
+ edge_feats.append(directions)
+ if edge_feats:
+ edge_feat["feat"] = np.concatenate(edge_feats, axis=-1).astype(
+ np.float32
+ )
+ if triplet_indices is not None:
+ edge_feat.update(triplet_indices)
+
+ return pgl.Graph(
+ num_nodes=positions.shape[0],
+ edges=edge_index.T.astype(np.int64),
+ node_feat=self.get_node_feat(atomic_numbers, positions),
+ edge_feat=edge_feat,
+ )
+
+ def get_triplet_indices(
+ self, edge_index: np.ndarray, num_atoms: int
+ ) -> Dict[str, np.ndarray]:
+ edge_index = np.asarray(edge_index, dtype=np.int64)
+ src, dst = edge_index
+ num_edges = edge_index.shape[1]
+
+ in_edges = [[] for _ in range(num_atoms)]
+ out_edges = [[] for _ in range(num_atoms)]
+ for edge_id in range(num_edges):
+ in_edges[int(dst[edge_id])].append(edge_id)
+ out_edges[int(src[edge_id])].append(edge_id)
+
+ idx_kj_list = []
+ idx_ji_list = []
+ for atom_id in range(num_atoms):
+ for kj_edge in in_edges[atom_id]:
+ k_atom = src[kj_edge]
+ for ji_edge in out_edges[atom_id]:
+ i_atom = dst[ji_edge]
+ if k_atom != i_atom:
+ idx_kj_list.append(kj_edge)
+ idx_ji_list.append(ji_edge)
+
+ idx_kj = np.asarray(idx_kj_list, dtype=np.int64)
+ idx_ji = np.asarray(idx_ji_list, dtype=np.int64)
+
+ idx_lk_list = []
+ idx_triplet_list = []
+ for triplet_id, kj_edge in enumerate(idx_kj_list):
+ k_atom = int(src[kj_edge])
+ lk_edges = in_edges[k_atom]
+ if lk_edges:
+ idx_lk_list.extend(lk_edges)
+ idx_triplet_list.extend([triplet_id] * len(lk_edges))
+
+ return {
+ "ti_idx_kj": idx_kj,
+ "ti_idx_ji": idx_ji,
+ "ti_idx_lk": np.asarray(idx_lk_list, dtype=np.int64),
+ "ti_idx_triplet": np.asarray(idx_triplet_list, dtype=np.int64),
+ }
+
+
def subgraph(
subset: Union[np.ndarray, List[int]],
edge_index: np.ndarray,
diff --git a/ppmat/models/common/initializer.py b/ppmat/models/common/initializer.py
index f5afa210..5b929192 100644
--- a/ppmat/models/common/initializer.py
+++ b/ppmat/models/common/initializer.py
@@ -37,6 +37,7 @@
"normal_",
"trunc_normal_",
"glorot_normal_",
+ "glorot_orthogonal_",
"constant_",
"ones_",
"zeros_",
@@ -490,3 +491,24 @@ def he_orthogonal_init(tensor):
tensor.data *= (1 / fan_in) ** 0.5
tensor.stop_gradient = stop_gradient
return tensor
+
+
+def glorot_orthogonal_(tensor: paddle.Tensor, scale: float = 1.0) -> paddle.Tensor:
+ """Orthogonal initialization with scaling (Glorot & Bengio, 2010 style).
+
+ Initializes the tensor as a random (semi-)orthogonal matrix and scales it
+ by the given factor. This is a simpler variant than ``he_orthogonal_init``
+ — it does not standardise or apply fan-in scaling.
+
+ Args:
+ tensor: Paddle Tensor or Parameter.
+ scale: Scaling factor applied after orthogonal initialisation.
+
+ Returns:
+ paddle.Tensor: The initialised tensor (same object, modified in-place).
+ """
+ init_orth = paddle.nn.initializer.Orthogonal()
+ init_orth(tensor)
+ with paddle.no_grad():
+ tensor.set_value(tensor * scale)
+ return tensor
diff --git a/ppmat/models/common/spherical_fourier_bessel.py b/ppmat/models/common/spherical_fourier_bessel.py
new file mode 100644
index 00000000..35223aeb
--- /dev/null
+++ b/ppmat/models/common/spherical_fourier_bessel.py
@@ -0,0 +1,340 @@
+# 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.
+"""Numerical spherical Fourier-Bessel embeddings for SphereNet."""
+
+import math
+from functools import lru_cache
+
+import numpy as np
+import paddle
+from scipy import optimize
+from scipy import special
+
+
+def _spherical_jn_root(x, order):
+ return special.spherical_jn(order, x)
+
+
+@lru_cache(maxsize=32)
+def _build_basis_constants(num_spherical, num_radial):
+ """Build deterministic Fourier-Bessel constants for one basis shape."""
+ if num_spherical < 1:
+ raise ValueError("num_spherical must be positive.")
+ if num_radial < 1:
+ raise ValueError("num_radial must be positive.")
+
+ zeros = np.zeros((num_spherical, num_radial), dtype=np.float64)
+ zeros[0] = np.arange(1, num_radial + 1, dtype=np.float64) * np.pi
+ points = np.arange(
+ 1, num_radial + num_spherical, dtype=np.float64
+ ) * np.pi
+ roots = np.zeros(num_radial + num_spherical - 1, dtype=np.float64)
+ for order in range(1, num_spherical):
+ for i in range(num_radial + num_spherical - 1 - order):
+ roots[i] = optimize.brentq(
+ _spherical_jn_root,
+ points[i],
+ points[i + 1],
+ args=(order,),
+ )
+ points = roots.copy()
+ zeros[order] = roots[:num_radial]
+
+ normalizers = np.empty_like(zeros)
+ for order in range(num_spherical):
+ values = special.spherical_jn(order + 1, zeros[order])
+ normalizers[order] = 1.0 / np.sqrt(0.5 * values**2)
+
+ harmonic_prefactors = np.zeros(
+ (num_spherical, num_spherical), dtype=np.float64
+ )
+ for degree in range(num_spherical):
+ for order in range(degree + 1):
+ harmonic_prefactors[degree, order] = math.sqrt(
+ (2 * degree + 1)
+ * math.factorial(degree - order)
+ / (4 * math.pi * math.factorial(degree + order))
+ )
+
+ zeros.setflags(write=False)
+ normalizers.setflags(write=False)
+ harmonic_prefactors.setflags(write=False)
+ return zeros, normalizers, harmonic_prefactors
+
+
+def _double_factorial(value):
+ result = 1
+ for factor in range(value, 0, -2):
+ result *= factor
+ return result
+
+
+def _spherical_jn_series(order, x, num_terms=20):
+ """Stable spherical-Bessel series around zero."""
+ x_squared = x * x
+ term = paddle.ones_like(x)
+ series = term
+ for index in range(1, num_terms):
+ term = (
+ term
+ * -x_squared
+ / (2 * index * (2 * order + 2 * index + 1))
+ )
+ series = series + term
+ return x**order / _double_factorial(2 * order + 1) * series
+
+
+class Envelope(paddle.nn.Layer):
+ """Smooth polynomial envelope function for radial cutoff."""
+
+ def __init__(self, exponent):
+ super().__init__()
+ self.p = exponent + 1
+ self.a = -(self.p + 1) * (self.p + 2) / 2
+ self.b = self.p * (self.p + 2)
+ self.c = -self.p * (self.p + 1) / 2
+
+ def forward(self, x):
+ p, a, b, c = self.p, self.a, self.b, self.c
+ x_pow_p0 = x.pow(p - 1)
+ x_pow_p1 = x_pow_p0 * x
+ x_pow_p2 = x_pow_p1 * x
+ return 1.0 / x + a * x_pow_p0 + b * x_pow_p1 + c * x_pow_p2
+
+
+class DistEmbedding(paddle.nn.Layer):
+ """Radial basis with a smooth envelope cutoff."""
+
+ def __init__(self, num_radial, cutoff=5.0, envelope_exponent=5):
+ super().__init__()
+ self.cutoff = cutoff
+ self.envelope = Envelope(envelope_exponent)
+ self.freq = paddle.create_parameter(
+ shape=[num_radial],
+ dtype=paddle.get_default_dtype(),
+ default_initializer=paddle.nn.initializer.Assign(
+ paddle.arange(
+ 1, num_radial + 1, dtype=paddle.get_default_dtype()
+ ).multiply(paddle.to_tensor(math.pi))
+ ),
+ )
+
+ def reset_parameters(self):
+ with paddle.no_grad():
+ self.freq.set_value(
+ paddle.arange(
+ 1,
+ self.freq.shape[0] + 1,
+ dtype=paddle.get_default_dtype(),
+ ).multiply(paddle.to_tensor(math.pi))
+ )
+
+ def forward(self, dist):
+ dist = dist.unsqueeze(-1) / self.cutoff
+ return self.envelope(dist) * paddle.sin(self.freq * dist)
+
+
+class SphericalBesselBasis(paddle.nn.Layer):
+ """Normalized spherical-Bessel basis evaluated with Paddle operations."""
+
+ def __init__(self, num_spherical, num_radial, cutoff=5.0):
+ super().__init__()
+ self.num_spherical = num_spherical
+ self.num_radial = num_radial
+ self.cutoff = cutoff
+
+ zeros, normalizers, _ = _build_basis_constants(
+ num_spherical, num_radial
+ )
+ self.register_buffer(
+ "zeros",
+ paddle.to_tensor(np.array(zeros, copy=True), dtype="float32"),
+ persistable=False,
+ )
+ self.register_buffer(
+ "normalizers",
+ paddle.to_tensor(
+ np.array(normalizers, copy=True), dtype="float32"
+ ),
+ persistable=False,
+ )
+
+ def forward(self, dist):
+ scaled_dist = dist.reshape([-1, 1, 1]) / self.cutoff
+ arguments = scaled_dist * self.zeros.reshape(
+ [1, self.num_spherical, self.num_radial]
+ )
+ safe_arguments = paddle.where(
+ paddle.abs(arguments) < 1e-2,
+ paddle.full_like(arguments, 1e-2),
+ arguments,
+ )
+
+ values = [paddle.sin(safe_arguments) / safe_arguments]
+ if self.num_spherical > 1:
+ values.append(
+ paddle.sin(safe_arguments) / safe_arguments**2
+ - paddle.cos(safe_arguments) / safe_arguments
+ )
+ for degree in range(1, self.num_spherical - 1):
+ values.append(
+ (2 * degree + 1) / safe_arguments * values[-1]
+ - values[-2]
+ )
+
+ basis = []
+ for degree in range(self.num_spherical):
+ degree_arguments = arguments[:, degree, :]
+ degree_values = values[degree][:, degree, :]
+ degree_values = paddle.where(
+ paddle.abs(degree_arguments) < degree + 1.0,
+ _spherical_jn_series(degree, degree_arguments),
+ degree_values,
+ )
+ basis.append(degree_values * self.normalizers[degree])
+ return paddle.stack(basis, axis=1)
+
+
+class RealSphericalHarmonics(paddle.nn.Layer):
+ """Real spherical harmonics with the SphereNet ordering convention."""
+
+ def __init__(self, num_spherical):
+ super().__init__()
+ self.num_spherical = num_spherical
+ _, _, prefactors = _build_basis_constants(num_spherical, 1)
+ self.register_buffer(
+ "prefactors",
+ paddle.to_tensor(
+ np.array(prefactors, copy=True), dtype="float32"
+ ),
+ persistable=False,
+ )
+ self.register_buffer(
+ "m0_indices",
+ paddle.to_tensor(
+ [degree * degree for degree in range(num_spherical)],
+ dtype="int64",
+ ),
+ persistable=False,
+ )
+
+ def forward(self, angle, torsion):
+ cos_angle = paddle.cos(angle)
+ sin_angle = paddle.sin(angle)
+ one = paddle.ones_like(cos_angle)
+ polynomials = {(0, 0): one}
+
+ for order in range(1, self.num_spherical):
+ polynomials[(order, order)] = (
+ 1 - 2 * order
+ ) * polynomials[(order - 1, order - 1)]
+
+ for order in range(self.num_spherical - 1):
+ polynomials[(order + 1, order)] = (
+ (2 * order + 1)
+ * cos_angle
+ * polynomials[(order, order)]
+ )
+
+ for order in range(self.num_spherical):
+ for degree in range(order + 2, self.num_spherical):
+ polynomials[(degree, order)] = (
+ (2 * degree - 1)
+ * cos_angle
+ * polynomials[(degree - 1, order)]
+ - (order + degree - 1)
+ * polynomials[(degree - 2, order)]
+ ) / (degree - order)
+
+ harmonics = []
+ sqrt_two = math.sqrt(2.0)
+ for degree in range(self.num_spherical):
+ harmonics.append(
+ self.prefactors[degree, 0]
+ * polynomials[(degree, 0)]
+ )
+ for order in range(1, degree + 1):
+ harmonics.append(
+ sqrt_two
+ * self.prefactors[degree, order]
+ * polynomials[(degree, order)]
+ * sin_angle**order
+ * paddle.cos(order * torsion)
+ )
+ for order in range(degree, 0, -1):
+ harmonics.append(
+ sqrt_two
+ * self.prefactors[degree, order]
+ * polynomials[(degree, order)]
+ * sin_angle**order
+ * paddle.sin(order * torsion)
+ )
+ return paddle.stack(harmonics, axis=1)
+
+
+class SphericalFourierBesselEmbedding(paddle.nn.Layer):
+ """Shared angle and torsion Fourier-Bessel embedding."""
+
+ def __init__(self, num_spherical, num_radial, cutoff=5.0):
+ super().__init__()
+ self.num_spherical = num_spherical
+ self.num_radial = num_radial
+ self.radial_basis = SphericalBesselBasis(
+ num_spherical, num_radial, cutoff
+ )
+ self.spherical_harmonics = RealSphericalHarmonics(num_spherical)
+
+ def forward(self, dist, angle, torsion, idx_kj):
+ radial_basis = self.radial_basis(dist)[idx_kj]
+ harmonics = self.spherical_harmonics(angle, torsion)
+ angle_harmonics = paddle.index_select(
+ harmonics,
+ self.spherical_harmonics.m0_indices,
+ axis=1,
+ )
+
+ num_triplets = angle.shape[0]
+ angle_embedding = (
+ radial_basis * angle_harmonics.reshape(
+ [num_triplets, self.num_spherical, 1]
+ )
+ ).reshape(
+ [num_triplets, self.num_spherical * self.num_radial]
+ )
+ torsion_embedding = (
+ radial_basis.reshape(
+ [
+ num_triplets,
+ 1,
+ self.num_spherical,
+ self.num_radial,
+ ]
+ )
+ * harmonics.reshape(
+ [
+ num_triplets,
+ self.num_spherical,
+ self.num_spherical,
+ 1,
+ ]
+ )
+ ).reshape(
+ [
+ num_triplets,
+ self.num_spherical
+ * self.num_spherical
+ * self.num_radial,
+ ]
+ )
+ return angle_embedding, torsion_embedding
diff --git a/ppmat/models/gmtnet/__init__.py b/ppmat/models/gmtnet/__init__.py
new file mode 100644
index 00000000..e5d64595
--- /dev/null
+++ b/ppmat/models/gmtnet/__init__.py
@@ -0,0 +1,3 @@
+from .gmtnet import GMTNet
+
+__all__ = ["GMTNet"]
diff --git a/ppmat/models/gmtnet/gmtnet.py b/ppmat/models/gmtnet/gmtnet.py
new file mode 100644
index 00000000..d1c2c861
--- /dev/null
+++ b/ppmat/models/gmtnet/gmtnet.py
@@ -0,0 +1,725 @@
+import copy
+import math
+from collections.abc import Mapping
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from ppmat.losses import build_loss
+from ppmat.models.common.e3nn import o3
+
+
+def get_arg(args, name, default):
+ if isinstance(args, dict):
+ return args.get(name, default)
+ return getattr(args, name, default)
+
+
+def pad_or_slice_last(x, target_dim):
+ """Pad or slice the last dimension to target_dim."""
+ current_dim = x.shape[-1]
+
+ if current_dim == target_dim:
+ return x
+
+ if current_dim > target_dim:
+ return x[..., :target_dim]
+
+ pad_shape = list(x.shape)
+ pad_shape[-1] = target_dim - current_dim
+ pad = paddle.zeros(pad_shape, dtype=x.dtype)
+
+ return paddle.concat([x, pad], axis=-1)
+
+
+def equality_adjustment(equality, batch):
+ """Paddle version of equality_adjustment.
+
+ This is mainly for eval/inference. It mirrors the PyTorch loop logic.
+ """
+ if equality is None:
+ return batch
+
+ out = batch.clone()
+ b, l1, l2 = out.shape
+ flat = out.reshape([b, l1 * l2])
+
+ for i in range(b):
+ mask = equality[i]
+ for j in range(l1 * l2):
+ for k in range(j + 1, l1 * l2):
+ if bool(mask[j, k].item()):
+ avg = (flat[i, j] + flat[i, k]) / 2.0
+ flat[i, j] = avg
+ flat[i, k] = avg
+
+ return flat.reshape([b, l1, l2])
+
+
+class SiLU(nn.Layer):
+ def forward(self, x):
+ return x * F.sigmoid(x)
+
+
+class RBFExpansion(nn.Layer):
+ """Paddle version of PyTorch RBFExpansion."""
+
+ def __init__(self, vmin=0.0, vmax=8.0, bins=40, lengthscale=None):
+ super().__init__()
+
+ self.vmin = vmin
+ self.vmax = vmax
+ self.bins = bins
+
+ centers = paddle.linspace(self.vmin, self.vmax, self.bins)
+ self.register_buffer("centers", centers)
+
+ if lengthscale is None:
+ # Match PyTorch:
+ # self.lengthscale = np.diff(self.centers).mean()
+ # self.gamma = 1 / self.lengthscale
+ if self.bins > 1:
+ self.lengthscale = float((self.vmax - self.vmin) / (self.bins - 1))
+ else:
+ self.lengthscale = float(self.vmax - self.vmin)
+
+ self.gamma = float(1.0 / self.lengthscale)
+ else:
+ self.lengthscale = float(lengthscale)
+ self.gamma = float(1.0 / (self.lengthscale**2))
+
+ def forward(self, distance):
+ return paddle.exp(-self.gamma * (distance.unsqueeze(1) - self.centers) ** 2)
+
+
+class ComformerConv(nn.Layer):
+ """First-stage Paddle implementation of ComformerConv.
+
+ This version focuses on matching module names and state_dict keys.
+ Full message passing logic will be aligned in the next stage.
+ """
+
+ def __init__(self, in_channels, out_channels, heads=1, edge_dim=None):
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.heads = heads
+ self.edge_dim = edge_dim if edge_dim is not None else out_channels
+
+ self.lin_key = nn.Linear(in_channels, out_channels)
+ self.lin_query = nn.Linear(in_channels, out_channels)
+ self.lin_value = nn.Linear(in_channels, out_channels)
+ self.lin_edge = nn.Linear(self.edge_dim, out_channels)
+ self.lin_concate = nn.Linear(out_channels, out_channels)
+
+ self.lin_msg_update = nn.Sequential(
+ nn.Linear(out_channels * 3, out_channels),
+ SiLU(),
+ nn.Linear(out_channels, out_channels),
+ )
+
+ self.key_update = nn.Sequential(
+ nn.Linear(out_channels * 3, out_channels),
+ SiLU(),
+ nn.Linear(out_channels, out_channels),
+ )
+
+ self.bn = nn.BatchNorm1D(out_channels)
+ self.bn_att = nn.BatchNorm1D(out_channels)
+ self.softplus = nn.Softplus()
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, x, edge_index=None, edge_attr=None):
+ """Paddle implementation of PyG MessagePassing forward.
+
+ PyG convention:
+ edge_index[0] = source node j
+ edge_index[1] = target node i
+ """
+
+ H, C = self.heads, self.out_channels
+
+ if isinstance(x, (tuple, list)):
+ x_src, x_dst = x
+ else:
+ x_src, x_dst = x, x
+
+ query = self.lin_query(x_dst).reshape([-1, H, C])
+ key = self.lin_key(x_src).reshape([-1, H, C])
+ value = self.lin_value(x_src).reshape([-1, H, C])
+
+ # Fallback path for temporary GMTNet.forward smoke tests.
+ # Real graph forward should pass edge_index and edge_attr.
+ if edge_index is None:
+ out = value.reshape([-1, H * C])
+ out = self.lin_concate(out)
+ return self.softplus(x_dst + out)
+
+ edge_index = edge_index.astype("int64")
+ src = edge_index[0]
+ dst = edge_index[1]
+
+ num_nodes = x_dst.shape[0]
+ num_edges = src.shape[0]
+
+ if edge_attr is None:
+ edge_attr = paddle.zeros([num_edges, self.edge_dim], dtype=x_dst.dtype)
+
+ query_i = paddle.gather(query, dst, axis=0)
+ key_i = paddle.gather(key, dst, axis=0)
+ key_j = paddle.gather(key, src, axis=0)
+
+ value_i = paddle.gather(value, dst, axis=0)
+ value_j = paddle.gather(value, src, axis=0)
+
+ edge_attr = self.lin_edge(edge_attr).reshape([-1, H, C])
+
+ key_j = self.key_update(paddle.concat([key_i, key_j, edge_attr], axis=-1))
+
+ alpha = (query_i * key_j) / (C**0.5)
+
+ out = self.lin_msg_update(paddle.concat([value_i, value_j, edge_attr], axis=-1))
+
+ alpha_bn = self.bn_att(alpha.reshape([-1, C])).reshape([-1, H, C])
+ out = out * self.sigmoid(alpha_bn)
+
+ # aggregate messages to target nodes with add aggregation
+ out_nodes = paddle.zeros([num_nodes, H, C], dtype=out.dtype)
+ index = dst.reshape([-1, 1])
+ out_nodes = paddle.scatter_nd_add(out_nodes, index, out)
+
+ out_nodes = out_nodes.reshape([-1, H * C])
+ out_nodes = self.lin_concate(out_nodes)
+
+ return self.softplus(x_dst + out_nodes)
+
+
+class W3JBuffers(nn.Layer):
+ """Container for e3nn Wigner-3j buffers.
+
+ This is a structure/key-compatible placeholder.
+ """
+
+ def __init__(self, shapes):
+ super().__init__()
+ for name, shape in shapes.items():
+ self.register_buffer(name, paddle.zeros(shape, dtype="float32"))
+
+
+class FakeTensorProduct(nn.Layer):
+ """Key-compatible placeholder for e3nn TensorProduct.
+
+ It registers the same state_dict keys used by the PyTorch checkpoint:
+ - weight
+ - output_mask
+ - _compiled_main_left_right._w3j_...
+ """
+
+ def __init__(self, output_mask_dim, w3j_shapes=None):
+ super().__init__()
+
+ self.register_buffer("weight", paddle.empty([0], dtype="float32"))
+ self.register_buffer(
+ "output_mask", paddle.ones([output_mask_dim], dtype="float32")
+ )
+
+ if w3j_shapes is not None and len(w3j_shapes) > 0:
+ self._compiled_main_left_right = self.add_sublayer(
+ "_compiled_main_left_right",
+ W3JBuffers(w3j_shapes),
+ )
+
+
+class TensorProductConvLayer(nn.Layer):
+ """Paddle e3nn implementation of TensorProductConvLayer."""
+
+ def __init__(
+ self,
+ in_irreps,
+ sh_irreps,
+ out_irreps,
+ n_edge_features,
+ residual=True,
+ ):
+ super().__init__()
+
+ self.in_irreps = in_irreps
+ self.out_irreps = out_irreps
+ self.sh_irreps = sh_irreps
+ self.residual = residual
+
+ self.tp = o3.FullyConnectedTensorProduct(
+ in_irreps,
+ sh_irreps,
+ out_irreps,
+ shared_weights=False,
+ )
+
+ self.fc = nn.Sequential(
+ nn.Linear(n_edge_features, n_edge_features),
+ nn.Softplus(),
+ nn.Linear(n_edge_features, self.tp.weight_numel),
+ )
+
+ def forward(
+ self, node_attr, edge_index, edge_attr, edge_sh, out_nodes=None, reduce="mean"
+ ):
+ edge_index = edge_index.astype("int64")
+
+ edge_src = edge_index[0]
+ edge_dst = edge_index[1]
+
+ if out_nodes is None:
+ out_nodes = node_attr.shape[0]
+
+ node_dst = paddle.gather(node_attr, edge_dst, axis=0)
+ weight = self.fc(edge_attr)
+
+ tp_out = self.tp(node_dst, edge_sh, weight)
+
+ out_dim = tp_out.shape[-1]
+
+ out = paddle.zeros([out_nodes, out_dim], dtype=tp_out.dtype)
+ out = paddle.scatter_nd_add(
+ out,
+ edge_src.reshape([-1, 1]),
+ tp_out,
+ )
+
+ if reduce == "mean":
+ counts = paddle.zeros([out_nodes, 1], dtype=tp_out.dtype)
+ counts = paddle.scatter_nd_add(
+ counts,
+ edge_src.reshape([-1, 1]),
+ paddle.ones([edge_src.shape[0], 1], dtype=tp_out.dtype),
+ )
+ out = out / paddle.clip(counts, min=1.0)
+
+ if self.residual:
+ padded = pad_or_slice_last(node_attr, out.shape[-1])
+ out = out + padded
+
+ return out
+
+
+class ComformerConvEqui(nn.Layer):
+ """Paddle e3nn implementation of ComformerConvEqui."""
+
+ def __init__(
+ self,
+ embsize=128,
+ ns=16,
+ nv=2,
+ residual=True,
+ ):
+ super().__init__()
+
+ irrep_seq = [
+ f"{ns}x0e",
+ f"{ns}x0e + {nv}x1o + {nv}x2e",
+ f"{ns}x0e + {nv}x1o + {nv}x1e + {nv}x2e + {nv}x2o",
+ "1x0e + 1x0o + 1x1e + 1x1o + 1x2e + 1x2o + 1x3e + 1x3o",
+ ]
+
+ self.ns = ns
+ self.nv = nv
+
+ self.node_linear = nn.Linear(embsize, ns)
+ self.sh = "1x0e + 1x1o + 1x2e"
+
+ self.nlayer_1 = TensorProductConvLayer(
+ in_irreps=irrep_seq[0],
+ sh_irreps=self.sh,
+ out_irreps=irrep_seq[1],
+ n_edge_features=embsize,
+ residual=residual,
+ )
+
+ self.nlayer_2 = TensorProductConvLayer(
+ in_irreps=irrep_seq[1],
+ sh_irreps=self.sh,
+ out_irreps=irrep_seq[2],
+ n_edge_features=embsize,
+ residual=False,
+ )
+
+ self.nlayer_3 = TensorProductConvLayer(
+ in_irreps=irrep_seq[2],
+ sh_irreps=self.sh,
+ out_irreps=irrep_seq[3],
+ n_edge_features=embsize,
+ residual=False,
+ )
+
+ def forward(self, data, node_features=None, edge_index=None, edge_features=None):
+ edge_vec = data.edge_attr
+
+ edge_irr = o3.spherical_harmonics(
+ self.sh,
+ edge_vec,
+ normalize=True,
+ normalization="component",
+ )
+
+ node_feature = self.node_linear(node_features)
+ node_feature = self.nlayer_1(node_feature, edge_index, edge_features, edge_irr)
+ node_feature = self.nlayer_2(node_feature, edge_index, edge_features, edge_irr)
+ node_feature = self.nlayer_3(node_feature, edge_index, edge_features, edge_irr)
+
+ return node_feature
+
+
+class GradientBlock(nn.Layer):
+ """Paddle e3nn eval implementation of Gradient_block.
+
+ Note:
+ create_graph=True is not supported because Paddle einsum_grad
+ currently does not support higher-order grad.
+ This implementation is for eval/test forward.
+ """
+
+ def __init__(self, nv=2):
+ super().__init__()
+
+ irrep_seq = [
+ "1x0e + 1x0o + 1x1e + 1x1o + 1x2e + 1x2o + 1x3e + 1x3o",
+ "1x1o",
+ ]
+
+ self.nv = nv
+ self.sh = "1x1o"
+
+ self.tp = o3.FullyConnectedTensorProduct(
+ irrep_seq[0],
+ self.sh,
+ irrep_seq[1],
+ internal_weights=False,
+ )
+
+ self.register_buffer(
+ "constant_w",
+ paddle.ones([self.tp.weight_numel], dtype="float32"),
+ )
+
+ def _training_dielectric(self, node_feature):
+ input_slices = self.tp.irreps_in1.slices()
+ field_mul_ir = self.tp.irreps_in2[0]
+ output_mul_ir = self.tp.irreps_out[0]
+
+ if field_mul_ir.mul != 1 or field_mul_ir.ir.l != 1:
+ raise RuntimeError("GradientBlock requires a single 1o field irrep.")
+ if output_mul_ir.mul != 1 or output_mul_ir.ir.l != 1:
+ raise RuntimeError("GradientBlock requires a single 1o output irrep.")
+
+ batch_size = node_feature.shape[0]
+ field_dim = field_mul_ir.ir.dim
+ output_dim = output_mul_ir.ir.dim
+
+ if field_dim != 3 or output_dim != 3:
+ raise RuntimeError("GradientBlock requires three-dimensional field/output.")
+
+ dielectric = paddle.zeros(
+ [batch_size, output_mul_ir.mul, output_dim, field_dim],
+ dtype=node_feature.dtype,
+ )
+ external_weight = self.constant_w.astype(node_feature.dtype)
+
+ for instruction_index, instruction, weight_view in self.tp.weight_views(
+ external_weight,
+ yield_instruction=True,
+ ):
+ if instruction.connection_mode != "uvw" or not instruction.has_weight:
+ raise RuntimeError(
+ "GradientBlock only supports weighted uvw tensor-product paths."
+ )
+ if instruction.i_in2 != 0 or instruction.i_out != 0:
+ raise RuntimeError(
+ "GradientBlock only supports the configured single field/output irrep."
+ )
+
+ input_mul_ir = self.tp.irreps_in1[instruction.i_in1]
+ path_output_mul_ir = self.tp.irreps_out[instruction.i_out]
+
+ if weight_view.shape != [
+ input_mul_ir.mul,
+ field_mul_ir.mul,
+ path_output_mul_ir.mul,
+ ]:
+ raise RuntimeError("Unexpected tensor-product external weight shape.")
+
+ input_feature = node_feature[:, input_slices[instruction.i_in1]]
+ input_feature = input_feature.reshape(
+ [batch_size, input_mul_ir.mul, input_mul_ir.ir.dim]
+ )
+
+ coupling = o3.wigner_3j(
+ input_mul_ir.ir.l,
+ field_mul_ir.ir.l,
+ path_output_mul_ir.ir.l,
+ dtype=node_feature.dtype,
+ device=node_feature.place,
+ )
+
+ feature_term = input_feature.reshape(
+ [batch_size, input_mul_ir.mul, input_mul_ir.ir.dim, 1, 1, 1]
+ )
+ coupling_term = coupling.reshape(
+ [1, 1, input_mul_ir.ir.dim, field_dim, output_dim, 1]
+ )
+ weight_term = weight_view[:, 0, :].reshape(
+ [1, input_mul_ir.mul, 1, 1, 1, path_output_mul_ir.mul]
+ )
+
+ coefficient_qkw = paddle.sum(
+ feature_term * coupling_term * weight_term,
+ axis=[1, 2],
+ )
+ coefficient_qkw = instruction.path_weight * coefficient_qkw
+ coefficient_wkq = coefficient_qkw.transpose([0, 3, 2, 1])
+
+ dielectric = dielectric + coefficient_wkq
+
+ spherical_derivative = math.sqrt(3.0 / (4.0 * math.pi))
+ dielectric = spherical_derivative * dielectric
+ return dielectric.reshape([batch_size, output_dim, field_dim])
+
+ def forward(self, node_feature):
+ if self.training:
+ return self._training_dielectric(node_feature)
+
+ bs = node_feature.shape[0]
+
+ outer_E = paddle.ones([bs, 3], dtype=node_feature.dtype)
+ outer_E.stop_gradient = False
+
+ E_ = o3.spherical_harmonics(
+ self.sh,
+ outer_E,
+ normalize=False,
+ )
+
+ D_ = self.tp(
+ node_feature,
+ E_,
+ self.constant_w.astype(node_feature.dtype),
+ )
+
+ dielectric = []
+
+ for i in range(3):
+ grad_outputs = paddle.zeros([bs, 3], dtype=node_feature.dtype)
+ grad_outputs[:, i] = 1.0
+
+ grad_i = paddle.grad(
+ outputs=[D_],
+ inputs=[outer_E],
+ grad_outputs=[grad_outputs],
+ create_graph=False,
+ retain_graph=True,
+ )[0]
+
+ dielectric.append(grad_i)
+
+ return paddle.stack(dielectric, axis=0).transpose([1, 0, 2])
+
+
+class GMTNet(nn.Layer):
+ """First-stage Paddle GMTNet with RBFExpansion and ComformerConv."""
+
+ requires_forward_grad = True
+
+ def __init__(self, args, loss_cfg=None):
+ super().__init__()
+
+ atom_input_features = get_arg(args, "atom_input_features", 92)
+ embsize = get_arg(args, "embedding_features", 128)
+ edge_features = get_arg(args, "edge_features", 512)
+ output_features = get_arg(args, "output_features", 9)
+ num_layers = get_arg(args, "num_layers", 2)
+
+ self.atom_embedding = nn.Linear(atom_input_features, embsize)
+
+ self.rbf = nn.Sequential(
+ RBFExpansion(vmin=-4.0, vmax=0.0, bins=edge_features),
+ nn.Linear(edge_features, embsize),
+ nn.Softplus(),
+ )
+
+ self.att_layers = nn.LayerList(
+ [
+ ComformerConv(
+ in_channels=embsize,
+ out_channels=embsize,
+ heads=1,
+ edge_dim=embsize,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.equi_update = ComformerConvEqui(embsize)
+ self.output_block = GradientBlock()
+
+ self.mask = get_arg(args, "use_mask", False)
+ self.reduce = get_arg(args, "reduce_cell", False)
+
+ self.etgnn_linear = nn.Linear(embsize, 1)
+
+ if loss_cfg is None:
+ self.loss_fn = None
+ else:
+ try:
+ self.loss_fn = build_loss(copy.deepcopy(loss_cfg))
+ except Exception as error:
+ raise ValueError(f"Invalid loss_cfg: {error}") from error
+
+ @staticmethod
+ def _validate_mapping_tensor(name, value, trailing_shape, dtype=None):
+ if not isinstance(value, paddle.Tensor):
+ raise TypeError(f"Mapping key '{name}' must be a Paddle Tensor")
+ if value.ndim != 3:
+ raise ValueError(
+ f"Mapping key '{name}' must have rank 3, got shape {list(value.shape)}"
+ )
+ if list(value.shape[1:]) != list(trailing_shape):
+ raise ValueError(
+ f"Mapping key '{name}' must have shape [B,{trailing_shape[0]},{trailing_shape[1]}], got {list(value.shape)}"
+ )
+ if dtype is not None and value.dtype != dtype:
+ raise TypeError(
+ f"Mapping key '{name}' must have dtype {dtype}, got {value.dtype}"
+ )
+
+ def _mapping_inputs(self, batch):
+ for key in ("graph", "feature_mask", "matrix_equal"):
+ if key not in batch:
+ raise ValueError(f"Mapping input is missing required key '{key}'")
+ feature_mask = batch["feature_mask"]
+ matrix_equal = batch["matrix_equal"]
+ self._validate_mapping_tensor("feature_mask", feature_mask, (32, 32))
+ self._validate_mapping_tensor(
+ "matrix_equal", matrix_equal, (9, 9), dtype=paddle.bool
+ )
+ if feature_mask.shape[0] != matrix_equal.shape[0]:
+ raise ValueError(
+ "Mapping feature_mask and matrix_equal batch sizes must match"
+ )
+ dielectric = batch.get("dielectric")
+ if dielectric is not None:
+ self._validate_mapping_tensor("dielectric", dielectric, (3, 3))
+ if dielectric.shape[0] != feature_mask.shape[0]:
+ raise ValueError(
+ "Mapping dielectric and feature_mask batch sizes must match"
+ )
+ return batch["graph"], feature_mask, matrix_equal, dielectric
+
+ def _forward_tensor(self, data, feat_mask, equality):
+ """Paddle GMTNet forward.
+
+ Current stage:
+ - atom_embedding: real
+ - RBF edge feature: real
+ - ComformerConv att_layers: real
+ - equi_update: placeholder
+ - output_block: placeholder
+
+ Therefore this is still not final metric reproduction.
+ """
+
+ node_features = self.atom_embedding(data.x)
+
+ edge_feat = -0.75 / paddle.norm(data.edge_attr, axis=1)
+ edge_features = self.rbf(edge_feat)
+
+ node_features = self.att_layers[0](
+ node_features,
+ data.edge_index,
+ edge_features,
+ )
+
+ node_features = self.att_layers[1](
+ node_features,
+ data.edge_index,
+ edge_features,
+ )
+
+ # Placeholder currently returns node_features unchanged.
+ node_features = self.equi_update(
+ data,
+ node_features,
+ data.edge_index,
+ edge_features,
+ )
+
+ # Manual global mean pooling by data.batch.
+ if hasattr(data, "batch"):
+ batch = data.batch.astype("int64")
+ else:
+ batch = paddle.zeros([node_features.shape[0]], dtype="int64")
+
+ num_graphs = int(paddle.max(batch).item()) + 1
+
+ crystal_features = paddle.zeros(
+ [num_graphs, node_features.shape[1]],
+ dtype=node_features.dtype,
+ )
+
+ crystal_features = paddle.scatter_nd_add(
+ crystal_features,
+ batch.reshape([-1, 1]),
+ node_features,
+ )
+
+ counts = paddle.zeros([num_graphs, 1], dtype=node_features.dtype)
+ counts = paddle.scatter_nd_add(
+ counts,
+ batch.reshape([-1, 1]),
+ paddle.ones([node_features.shape[0], 1], dtype=node_features.dtype),
+ )
+
+ crystal_features = crystal_features / paddle.clip(counts, min=1.0)
+
+ if self.mask and feat_mask is not None:
+ crystal_features = paddle.bmm(
+ feat_mask,
+ crystal_features.unsqueeze(-1),
+ ).squeeze(-1)
+
+ outputs = self.output_block(crystal_features)
+
+ if equality is not None:
+ outputs = equality_adjustment(equality, outputs)
+
+ return outputs
+
+ def forward(self, data, feat_mask=None, equality=None):
+ if isinstance(data, Mapping):
+ if feat_mask is not None or equality is not None:
+ raise ValueError(
+ "Mapping input must not be combined with feat_mask or equality arguments"
+ )
+ graph, feature_mask, matrix_equal, dielectric = self._mapping_inputs(data)
+ prediction = self._forward_tensor(graph, feature_mask, matrix_equal)
+ result = {"pred_dict": {"dielectric": prediction}}
+ if dielectric is None:
+ return result
+ if self.loss_fn is None:
+ raise ValueError(
+ "Dielectric labels were provided, but no loss function was explicitly configured. No default loss is available."
+ )
+ result["loss_dict"] = {"loss": self.loss_fn(prediction, dielectric)}
+ return result
+ if feat_mask is None or equality is None:
+ raise ValueError(
+ "Raw GMTNet calls require both feat_mask and equality arguments"
+ )
+ return self._forward_tensor(data, feat_mask, equality)
+
+ def predict(self, data):
+ """Predict dielectric tensors from label-free GMTNet Mapping inputs."""
+ if isinstance(data, list):
+ predictions = [self.forward(item)["pred_dict"]["dielectric"] for item in data]
+ return {"pred_dict": {"dielectric": paddle.concat(predictions, axis=0)}}
+ return self.forward(data)
diff --git a/ppmat/models/gmtnet/gmtnet_graph_converter.py b/ppmat/models/gmtnet/gmtnet_graph_converter.py
new file mode 100644
index 00000000..3bcb6792
--- /dev/null
+++ b/ppmat/models/gmtnet/gmtnet_graph_converter.py
@@ -0,0 +1,305 @@
+"""Historical-equivalent crystal graph conversion for GMTNet.
+
+This converter reproduces the graph-construction portion of the historical
+GMTNet ``atoms2graphs`` path for ordered pymatgen structures. It converts the
+structure to JARVIS ``Atoms``, constructs canonicalized periodic neighbors,
+and returns a PaddleMaterials geometric ``Data`` object containing ``x``,
+``edge_index``, and ``edge_attr``.
+
+The first version intentionally serves only the frozen GMTNet graph contract.
+Callers must provide ``equivalent_atoms`` with one entry per structure atom.
+It does not generate feature masks, matrix-equality constraints, labels, data
+splits, or batches.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Mapping
+from collections.abc import Sequence
+
+import numpy as np
+import paddle
+import spglib
+from jarvis.core.specie import get_node_attributes
+from pymatgen.core import Structure
+from pymatgen.io.jarvis import JarvisAtomsAdaptor
+
+from ppmat.datasets.geometric_data_type.data import Data
+from ppmat.models.common.e3nn import o3
+from ppmat.models.common.e3nn.io import CartesianTensor
+
+
+class GMTNetGraphConverter:
+ """Convert an ordered pymatgen structure into the GMTNet graph format.
+
+ Args:
+ cutoff: Initial periodic-neighbor cutoff in Angstrom.
+ max_neighbors: Minimum number of neighbors used to select each site.
+ atom_features: JARVIS node-feature representation. Only ``"cgcnn"``
+ is supported by this historical-equivalence implementation.
+ use_canonize: Whether to canonicalize periodic edge representations.
+ reduce_cell: Historical ``reduce`` argument. The current GMTNet
+ baseline requires ``False``.
+
+ The returned :class:`~ppmat.datasets.geometric_data_type.data.Data` has
+ ``x`` with shape ``[num_nodes, 92]``, ``edge_index`` with shape
+ ``[2, num_edges]``, and ``edge_attr`` with shape ``[num_edges, 3]``.
+ """
+
+ def __init__(
+ self,
+ cutoff: float = 4.0,
+ max_neighbors: int = 16,
+ atom_features: str = "cgcnn",
+ use_canonize: bool = True,
+ reduce_cell: bool = False,
+ ):
+ if atom_features != "cgcnn":
+ raise ValueError(
+ "GMTNetGraphConverter only supports atom_features='cgcnn'."
+ )
+ if cutoff <= 0:
+ raise ValueError("cutoff must be positive.")
+ if max_neighbors <= 0:
+ raise ValueError("max_neighbors must be positive.")
+ if reduce_cell:
+ raise ValueError(
+ "GMTNetGraphConverter v1 supports only reduce_cell=False."
+ )
+
+ self.cutoff = float(cutoff)
+ self.max_neighbors = int(max_neighbors)
+ self.atom_features = atom_features
+ self.use_canonize = bool(use_canonize)
+ self.reduce_cell = bool(reduce_cell)
+ self._adaptor = JarvisAtomsAdaptor()
+ self._irreps_output = o3.Irreps(
+ "1x0e + 1x0o + 1x1e + 1x1o + 1x2e + 1x2o + 1x3e + 1x3o"
+ )
+ self._cartesian_converter = CartesianTensor("ij")
+ self._symprec = 1e-5
+
+ @staticmethod
+ def load_structure_from_cif(cif_path) -> Structure:
+ """Load a CIF without parser-side coordinate idealization for GMTNet."""
+ try:
+ return Structure.from_file(
+ cif_path,
+ primitive=False,
+ sort=False,
+ merge_tol=0.0,
+ frac_tolerance=0.0,
+ )
+ except Exception as error:
+ raise ValueError(
+ f"GMTNet precision-preserving CIF parsing failed for {cif_path!s}."
+ ) from error
+
+ @staticmethod
+ def _dataset_value(dataset, name: str):
+ try:
+ return getattr(dataset, name)
+ except AttributeError:
+ return dataset[name]
+
+ @staticmethod
+ def _unique_rotations(rotations: np.ndarray) -> np.ndarray:
+ unique_rotations = []
+ seen = set()
+ for rotation in rotations:
+ key = tuple(rotation.reshape(-1).tolist())
+ if key not in seen:
+ seen.add(key)
+ unique_rotations.append(rotation)
+ return np.asarray(unique_rotations, dtype=np.float32)
+
+ def _symmetry_dataset(self, structure: Structure):
+ dataset = spglib.get_symmetry_dataset(
+ (structure.lattice.matrix, structure.frac_coords, structure.atomic_numbers),
+ symprec=self._symprec,
+ )
+ if dataset is None:
+ raise ValueError("spglib could not determine the structure symmetry.")
+ return dataset
+
+ def _prediction_constraints(self, structure: Structure, symmetry_dataset):
+ rotations = self._unique_rotations(
+ np.asarray(self._dataset_value(symmetry_dataset, "rotations"))
+ )
+ lattice = np.asarray(structure.lattice.matrix, dtype=np.float32).T
+ transformed_rotations = lattice @ rotations @ np.linalg.inv(lattice)
+ representations = self._irreps_output.D_from_matrix(
+ paddle.to_tensor(transformed_rotations, dtype="float32")
+ )
+ average = representations.sum(axis=0) / representations.shape[0]
+ feature_mask = average * (average > 1e-5).astype("float32")
+ mask = paddle.concat([
+ paddle.arange(8, dtype="float32") + 10.0,
+ (paddle.arange(24, dtype="float32") + 18.0) * 100.0,
+ ])
+ feature_total = representations.sum(axis=0) @ mask
+ selected = feature_total[[0, 2, 3, 4, 8, 9, 10, 11, 12]]
+ ideal_matrix = self._cartesian_converter.to_cartesian(selected).numpy()
+ flattened = ideal_matrix.reshape(-1)
+ matrix_equal = np.abs(flattened[:, None] - flattened[None, :]) < (
+ 0.0001 * np.abs(flattened[:, None] + flattened[None, :]) / 2.0
+ )
+ return feature_mask, matrix_equal
+
+ def _build_prediction_input_from_structure(self, structure: Structure):
+ symmetry_dataset = self._symmetry_dataset(structure)
+ equivalent_atoms = np.asarray(
+ self._dataset_value(symmetry_dataset, "equivalent_atoms"), dtype=np.int32
+ )
+ feature_mask, matrix_equal = self._prediction_constraints(
+ structure, symmetry_dataset
+ )
+ return {
+ "graph": self(structure, equivalent_atoms),
+ "feature_mask": feature_mask.unsqueeze(0),
+ "matrix_equal": paddle.to_tensor(matrix_equal, dtype="bool").unsqueeze(0),
+ }
+
+ def build_prediction_input(self, values):
+ """Return GMTNet prediction Mappings while preserving Mapping inputs."""
+ if isinstance(values, Mapping):
+ return values
+ if isinstance(values, Structure):
+ return self._build_prediction_input_from_structure(values)
+ if not isinstance(values, Sequence) or isinstance(values, (str, bytes)):
+ raise TypeError("structures must be a pymatgen Structure or a sequence of Structures.")
+ if not values:
+ raise ValueError("structures must not be empty.")
+ if all(isinstance(value, Mapping) for value in values):
+ return values
+ if not all(isinstance(value, Structure) for value in values):
+ raise TypeError("structures must contain only pymatgen Structure objects.")
+ return [self._build_prediction_input_from_structure(value) for value in values]
+
+ @staticmethod
+ def _canonize_edge(src_id, dst_id, src_image, dst_image):
+ if dst_id < src_id:
+ src_id, dst_id = dst_id, src_id
+ src_image, dst_image = dst_image, src_image
+
+ if not np.array_equal(src_image, (0, 0, 0)):
+ shift = src_image
+ src_image = tuple(np.subtract(src_image, shift))
+ dst_image = tuple(np.subtract(dst_image, shift))
+
+ if src_image != (0, 0, 0):
+ raise RuntimeError("Canonical edge source image must be (0, 0, 0).")
+ return src_id, dst_id, src_image, dst_image
+
+ def _build_edges(self, atoms, cutoff: float):
+ all_neighbors = atoms.get_all_neighbors(r=cutoff)
+ if len(all_neighbors) == 0:
+ raise RuntimeError("Structure contains no atoms.")
+
+ min_neighbors = min(len(neighbor_list) for neighbor_list in all_neighbors)
+ if min_neighbors < self.max_neighbors:
+ lattice = atoms.lattice
+ expanded_cutoff = (
+ max(lattice.a, lattice.b, lattice.c)
+ if cutoff < max(lattice.a, lattice.b, lattice.c)
+ else 2.0 * cutoff
+ )
+ return self._build_edges(atoms, expanded_cutoff)
+
+ edges = defaultdict(set)
+ for site_index, neighbor_list in enumerate(all_neighbors):
+ neighbor_list = sorted(neighbor_list, key=lambda neighbor: neighbor[2])
+ distances = np.asarray([neighbor[2] for neighbor in neighbor_list])
+ neighbor_ids = np.asarray([neighbor[1] for neighbor in neighbor_list])
+ images = np.asarray([neighbor[3] for neighbor in neighbor_list])
+
+ max_distance = distances[self.max_neighbors - 1]
+ selected = distances <= max_distance
+ neighbor_ids = neighbor_ids[selected]
+ images = images[selected]
+
+ for destination_index, image in zip(neighbor_ids, images):
+ source_index, destination_index, _, destination_image = (
+ self._canonize_edge(
+ site_index,
+ int(destination_index),
+ (0, 0, 0),
+ tuple(image),
+ )
+ )
+ if self.use_canonize:
+ edges[(source_index, destination_index)].add(destination_image)
+ else:
+ edges[(site_index, int(destination_index))].add(tuple(image))
+ return edges
+
+ @staticmethod
+ def _build_undirected_edge_data(atoms, edges):
+ sources = []
+ destinations = []
+ displacements = []
+ for (source_index, destination_index), images in edges.items():
+ for destination_image in images:
+ destination_coordinate = (
+ atoms.frac_coords[destination_index] + destination_image
+ )
+ displacement = atoms.lattice.cart_coords(
+ destination_coordinate - atoms.frac_coords[source_index]
+ )
+ for source, destination, vector in (
+ (source_index, destination_index, displacement),
+ (destination_index, source_index, -displacement),
+ ):
+ sources.append(source)
+ destinations.append(destination)
+ displacements.append(vector)
+
+ if not sources:
+ raise RuntimeError("Graph construction produced no edges.")
+
+ edge_index = np.asarray([sources, destinations], dtype=np.int64)
+ edge_attr = np.asarray(displacements, dtype=np.float32)
+ return edge_index, edge_attr
+
+ def __call__(self, structure, equivalent_atoms: Sequence[int]):
+ """Return the historical-equivalent GMTNet graph for ``structure``.
+
+ Args:
+ structure: Ordered pymatgen ``Structure`` accepted by
+ :class:`pymatgen.io.jarvis.JarvisAtomsAdaptor`.
+ equivalent_atoms: Per-atom symmetry-equivalence identifiers. They
+ are required to preserve the historical call contract even
+ though the frozen ``reduce_cell=False`` path does not reduce
+ nodes by equivalence classes.
+ """
+ if structure is None:
+ raise ValueError("structure must be provided.")
+ if equivalent_atoms is None:
+ raise ValueError("equivalent_atoms must be provided.")
+
+ atoms = self._adaptor.get_atoms(structure)
+ if len(equivalent_atoms) != len(atoms.elements):
+ raise ValueError(
+ "equivalent_atoms length must equal the number of structure atoms."
+ )
+
+ edges = self._build_edges(atoms, self.cutoff)
+ edge_index, edge_attr = self._build_undirected_edge_data(atoms, edges)
+
+ node_features = np.asarray(
+ [
+ list(get_node_attributes(element, atom_features=self.atom_features))
+ for element in atoms.elements
+ ],
+ dtype=np.float32,
+ )
+ graph = Data(
+ x=paddle.to_tensor(node_features, dtype="float32"),
+ edge_index=paddle.to_tensor(edge_index, dtype="int64"),
+ edge_attr=paddle.to_tensor(edge_attr, dtype="float32"),
+ )
+ required_fields = ("x", "edge_index", "edge_attr")
+ if any(getattr(graph, field, None) is None for field in required_fields):
+ raise RuntimeError("GMTNet graph is missing required fields.")
+ return graph
diff --git a/ppmat/models/spherenet/geometry.py b/ppmat/models/spherenet/geometry.py
new file mode 100644
index 00000000..35cc2d0f
--- /dev/null
+++ b/ppmat/models/spherenet/geometry.py
@@ -0,0 +1,80 @@
+# 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.
+"""
+3D geometry utilities for spherical message-passing models.
+
+SphereNet-specific — computes distance, angle, and torsion from 3D
+atomic coordinates for spherical message passing.
+"""
+
+import paddle
+
+from ppmat.utils.scatter import scatter_argmin
+
+
+def compute_geometry(pos, edge_index, triplet_indices):
+ """Compute SphereNet distance, angle, and torsion tensors.
+
+ Args:
+ pos: Atom positions with shape [num_nodes, 3].
+ edge_index: Directed edge indices with shape [2, num_edges].
+ triplet_indices: Precomputed edge and triplet index tensors.
+
+ Returns:
+ Distance, angle, torsion, edge source/target indices, and triplet maps.
+ """
+ i, j = edge_index
+ idx_kj = triplet_indices["idx_kj"]
+ idx_ji = triplet_indices["idx_ji"]
+ idx_lk = triplet_indices["idx_lk"]
+ idx_triplet = triplet_indices["idx_triplet"]
+
+ vec = pos[j] - pos[i]
+ dist = paddle.sqrt(paddle.sum(vec * vec, axis=-1) + 1e-8)
+
+ vec_kj = vec[idx_kj]
+ vec_ji = vec[idx_ji]
+
+ angle_cross = paddle.linalg.cross(vec_kj, vec_ji)
+ angle_sin = paddle.sqrt(paddle.sum(angle_cross * angle_cross, axis=-1) + 1e-8)
+ angle_cos = -paddle.sum(vec_kj * vec_ji, axis=-1)
+ angle = paddle.atan2(angle_sin, angle_cos).detach()
+
+ torsion = paddle.zeros_like(angle)
+ if idx_lk.shape[0] > 0:
+ k_idx_from_edge_lk = j[idx_lk]
+ v1 = pos[k_idx_from_edge_lk] - pos[i[idx_lk]]
+ v2 = vec_kj[idx_triplet]
+ v3 = vec_ji[idx_triplet]
+
+ v2_norm = paddle.sqrt(paddle.sum(v2 * v2, axis=-1) + 1e-8)
+ v2_cross_v3 = paddle.linalg.cross(v2, v3)
+ v1_dot_v2crossv3 = paddle.sum(v1 * v2_cross_v3, axis=-1)
+ v1_cross_v2 = paddle.linalg.cross(v1, v2)
+ v1crossv2_dot_v2crossv3 = paddle.sum(v1_cross_v2 * v2_cross_v3, axis=-1)
+
+ torsion_angle = paddle.atan2(
+ v2_norm * v1_dot_v2crossv3, v1crossv2_dot_v2crossv3
+ ).detach()
+
+ torsion_indices = scatter_argmin(
+ paddle.abs(torsion_angle), idx_triplet, idx_kj.shape[0]
+ )
+ torsion = paddle.where(
+ torsion_indices >= 0,
+ torsion_angle[paddle.clip(torsion_indices, min=0)],
+ torsion,
+ )
+
+ return dist, angle, torsion, i, j, idx_kj, idx_ji
diff --git a/ppmat/models/spherenet/spherenet.py b/ppmat/models/spherenet/spherenet.py
new file mode 100644
index 00000000..8062c5b7
--- /dev/null
+++ b/ppmat/models/spherenet/spherenet.py
@@ -0,0 +1,452 @@
+# 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 paddle import nn
+from paddle.nn import Embedding
+from paddle.nn import Linear
+from paddle.nn.functional import silu
+
+from ppmat.models.common import initializer
+from ppmat.models.common.spherical_fourier_bessel import DistEmbedding
+from ppmat.models.common.spherical_fourier_bessel import (
+ SphericalFourierBesselEmbedding,
+)
+from ppmat.models.spherenet.geometry import compute_geometry
+from ppmat.utils.scatter import scatter_sum
+
+
+class SphereNetEmbedding(paddle.nn.Layer):
+ def __init__(self, num_spherical, num_radial, cutoff, envelope_exponent):
+ super().__init__()
+ self.dist_emb = DistEmbedding(num_radial, cutoff, envelope_exponent)
+ self.geometry_emb = SphericalFourierBesselEmbedding(
+ num_spherical, num_radial, cutoff
+ )
+
+ def reset_parameters(self):
+ self.dist_emb.reset_parameters()
+
+ def forward(self, dist, angle, torsion, idx_kj):
+ dist_emb = self.dist_emb(dist)
+ angle_emb, torsion_emb = self.geometry_emb(
+ dist, angle, torsion, idx_kj
+ )
+ return dist_emb, angle_emb, torsion_emb
+
+
+class ResidualLayer(paddle.nn.Layer):
+ def __init__(self, hidden_channels, act=silu):
+ super().__init__()
+ self.act = act
+ self.lin1 = Linear(hidden_channels, hidden_channels)
+ self.lin2 = Linear(hidden_channels, hidden_channels)
+
+ def reset_parameters(self):
+ initializer.glorot_orthogonal_(self.lin1.weight, scale=1.0)
+ initializer.zeros_(self.lin1.bias)
+ initializer.glorot_orthogonal_(self.lin2.weight, scale=1.0)
+ initializer.zeros_(self.lin2.bias)
+
+ def forward(self, x):
+ return x + self.act(self.lin2(self.act(self.lin1(x))))
+
+
+class InitialEdgeEmbedding(paddle.nn.Layer):
+ def __init__(
+ self,
+ num_radial,
+ hidden_channels,
+ act=silu,
+ use_node_features=True,
+ use_extra_node_feature=False,
+ ):
+ super().__init__()
+ self.act = act
+ self.use_node_features = use_node_features
+ self.use_extra_node_feature = use_extra_node_feature
+ if self.use_node_features:
+ self.emb = Embedding(95, hidden_channels)
+ else:
+ self.node_embedding = paddle.create_parameter(
+ shape=[hidden_channels],
+ dtype=paddle.get_default_dtype(),
+ default_initializer=paddle.nn.initializer.Normal(),
+ )
+ self.lin_rbf_0 = Linear(num_radial, hidden_channels)
+ if self.use_extra_node_feature:
+ self.lin = Linear(5 * hidden_channels, hidden_channels)
+ else:
+ self.lin = Linear(3 * hidden_channels, hidden_channels)
+ self.lin_rbf_1 = Linear(num_radial, hidden_channels, bias_attr=False)
+
+ def reset_parameters(self):
+ if self.use_node_features:
+ initializer.uniform_(
+ self.emb.weight,
+ -(3.0**0.5),
+ 3.0**0.5,
+ )
+ else:
+ initializer.normal_(self.node_embedding)
+
+ for layer in self.sublayers():
+ if isinstance(layer, Linear):
+ initializer.glorot_orthogonal_(layer.weight, scale=1.0)
+ if layer.bias is not None:
+ initializer.zeros_(layer.bias)
+
+ def forward(self, x, node_feature, emb_in, i, j):
+ rbf, _, _ = emb_in
+ if self.use_node_features:
+ x = self.emb(x)
+ else:
+ x = self.node_embedding.unsqueeze(0).expand([x.shape[0], -1])
+ if node_feature is not None and self.use_extra_node_feature:
+ x = paddle.concat([x, node_feature], axis=1)
+ rbf0 = self.act(self.lin_rbf_0(rbf))
+ e1 = self.act(self.lin(paddle.concat([x[i], x[j], rbf0], axis=-1)))
+ e2 = self.lin_rbf_1(rbf) * e1
+ return e1, e2
+
+
+class EdgeUpdate(paddle.nn.Layer):
+ def __init__(
+ self,
+ hidden_channels,
+ int_emb_size,
+ basis_emb_size_dist,
+ basis_emb_size_angle,
+ basis_emb_size_torsion,
+ num_spherical,
+ num_radial,
+ num_before_skip,
+ num_after_skip,
+ act=silu,
+ ):
+ super().__init__()
+ self.act = act
+ self.lin_rbf1 = Linear(num_radial, basis_emb_size_dist, bias_attr=False)
+ self.lin_rbf2 = Linear(basis_emb_size_dist, hidden_channels, bias_attr=False)
+ self.lin_sbf1 = Linear(
+ num_spherical * num_radial, basis_emb_size_angle, bias_attr=False
+ )
+ self.lin_sbf2 = Linear(basis_emb_size_angle, int_emb_size, bias_attr=False)
+ self.lin_t1 = Linear(
+ num_spherical * num_spherical * num_radial,
+ basis_emb_size_torsion,
+ bias_attr=False,
+ )
+ self.lin_t2 = Linear(basis_emb_size_torsion, int_emb_size, bias_attr=False)
+ self.lin_rbf = Linear(num_radial, hidden_channels, bias_attr=False)
+
+ self.lin_kj = Linear(hidden_channels, hidden_channels)
+ self.lin_ji = Linear(hidden_channels, hidden_channels)
+
+ self.lin_down = Linear(hidden_channels, int_emb_size, bias_attr=False)
+ self.lin_up = Linear(int_emb_size, hidden_channels, bias_attr=False)
+
+ self.layers_before_skip = nn.LayerList(
+ [ResidualLayer(hidden_channels, act) for _ in range(num_before_skip)]
+ )
+ self.lin = Linear(hidden_channels, hidden_channels)
+ self.layers_after_skip = nn.LayerList(
+ [ResidualLayer(hidden_channels, act) for _ in range(num_after_skip)]
+ )
+
+ def reset_parameters(self):
+ for layer in self.sublayers():
+ if isinstance(layer, Linear):
+ initializer.glorot_orthogonal_(layer.weight, scale=1.0)
+ if layer.bias is not None:
+ initializer.zeros_(layer.bias)
+
+ def forward(self, x, emb_in, idx_kj, idx_ji):
+ rbf0, sbf, t = emb_in
+ x1, _ = x
+
+ x_ji = self.act(self.lin_ji(x1))
+ x_kj = self.act(self.lin_kj(x1))
+
+ rbf = self.lin_rbf1(rbf0)
+ rbf = self.lin_rbf2(rbf)
+ x_kj = x_kj * rbf
+
+ x_kj = self.act(self.lin_down(x_kj))
+
+ sbf = self.lin_sbf1(sbf)
+ sbf = self.lin_sbf2(sbf)
+ x_kj = x_kj[idx_kj] * sbf
+
+ t = self.lin_t1(t)
+ t = self.lin_t2(t)
+ x_kj = x_kj * t
+
+ x_kj = scatter_sum(x_kj, idx_ji, dim=0, dim_size=x1.shape[0])
+ x_kj = self.act(self.lin_up(x_kj))
+
+ e1 = x_ji + x_kj
+ for layer in self.layers_before_skip:
+ e1 = layer(e1)
+ e1 = self.act(self.lin(e1)) + x1
+ for layer in self.layers_after_skip:
+ e1 = layer(e1)
+ e2 = self.lin_rbf(rbf0) * e1
+ return e1, e2
+
+
+class NodeUpdate(paddle.nn.Layer):
+ def __init__(
+ self,
+ hidden_channels,
+ out_emb_channels,
+ out_channels,
+ num_output_layers,
+ act,
+ output_init,
+ ):
+ super().__init__()
+ self.act = act
+ self.output_init = output_init
+
+ self.lin_up = Linear(hidden_channels, out_emb_channels, bias_attr=True)
+ self.lins = nn.LayerList()
+ for _ in range(num_output_layers):
+ self.lins.append(Linear(out_emb_channels, out_emb_channels))
+ self.lin = Linear(out_emb_channels, out_channels, bias_attr=False)
+
+ def reset_parameters(self):
+ initializer.glorot_orthogonal_(self.lin_up.weight, scale=1.0)
+ for lin in self.lins:
+ initializer.glorot_orthogonal_(lin.weight, scale=1.0)
+ initializer.zeros_(lin.bias)
+ if self.output_init == "zeros":
+ initializer.zeros_(self.lin.weight)
+ if self.output_init == "GlorotOrthogonal":
+ initializer.glorot_orthogonal_(self.lin.weight, scale=1.0)
+
+ def forward(self, e, i, dim_size=None):
+ _, e2 = e
+ v = scatter_sum(e2, i, dim=0, dim_size=dim_size)
+ v = self.lin_up(v)
+ for lin in self.lins:
+ v = self.act(lin(v))
+ v = self.lin(v)
+ return v
+
+
+class SphereNet(paddle.nn.Layer):
+ """Spherical Message Passing for 3D molecular graph tasks.
+
+ This class follows the PaddleMaterials model protocol directly: ``forward``
+ accepts a batch dict and returns ``loss_dict`` / ``pred_dict``. Core
+ tensor computation is handled by ``_forward``.
+ """
+
+ def __init__(
+ self,
+ energy_and_force=False,
+ cutoff=5.0,
+ num_layers=4,
+ hidden_channels=128,
+ out_channels=1,
+ int_emb_size=64,
+ basis_emb_size_dist=8,
+ basis_emb_size_angle=8,
+ basis_emb_size_torsion=8,
+ out_emb_channels=256,
+ num_spherical=7,
+ num_radial=6,
+ envelope_exponent=5,
+ num_before_skip=1,
+ num_after_skip=2,
+ num_output_layers=3,
+ act="swish",
+ output_init="GlorotOrthogonal",
+ use_node_features=True,
+ use_extra_node_feature=False,
+ extra_node_feature_dim=1,
+ property_name="mu",
+ force_key="force",
+ ):
+ super().__init__()
+
+ act_fn = silu if act in ("swish", "silu") else act
+ if not callable(act_fn):
+ raise ValueError(f"Unsupported activation: {act}")
+
+ self.energy_and_force = energy_and_force
+ self.use_extra_node_feature = use_extra_node_feature
+ self.property_name = property_name
+ self.force_key = force_key
+
+ if use_extra_node_feature:
+ self.extra_emb = Linear(extra_node_feature_dim, hidden_channels)
+
+ self.init_e = InitialEdgeEmbedding(
+ num_radial,
+ hidden_channels,
+ act_fn,
+ use_node_features=use_node_features,
+ use_extra_node_feature=use_extra_node_feature,
+ )
+ node_update_cfg = {
+ "hidden_channels": hidden_channels,
+ "out_emb_channels": out_emb_channels,
+ "out_channels": out_channels,
+ "num_output_layers": num_output_layers,
+ "act": act_fn,
+ "output_init": output_init,
+ }
+ self.init_v = NodeUpdate(**node_update_cfg)
+ self.emb_layer = SphereNetEmbedding(
+ num_spherical, num_radial, cutoff, envelope_exponent
+ )
+
+ self.update_vs = nn.LayerList(
+ [NodeUpdate(**node_update_cfg) for _ in range(num_layers)]
+ )
+
+ self.update_es = nn.LayerList(
+ [
+ EdgeUpdate(
+ hidden_channels,
+ int_emb_size,
+ basis_emb_size_dist,
+ basis_emb_size_angle,
+ basis_emb_size_torsion,
+ num_spherical,
+ num_radial,
+ num_before_skip,
+ num_after_skip,
+ act_fn,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.reset_parameters()
+
+ def reset_parameters(self):
+ if self.use_extra_node_feature:
+ initializer.glorot_orthogonal_(self.extra_emb.weight, scale=1.0)
+ initializer.zeros_(self.extra_emb.bias)
+ layers = [
+ self.init_e,
+ self.init_v,
+ self.emb_layer,
+ *self.update_es,
+ *self.update_vs,
+ ]
+ for layer in layers:
+ layer.reset_parameters()
+
+ def _forward(self, data):
+ graph = data["graph"].tensor()
+ z = graph.node_feat["atomic_number"].astype("int64").reshape([-1])
+ pos = graph.node_feat["pos"].astype(paddle.get_default_dtype())
+ if self.energy_and_force:
+ pos = pos.detach()
+ pos.stop_gradient = False
+
+ node_batch = graph.graph_node_id.astype("int64")
+ edge_index = paddle.transpose(graph.edges.astype("int64"), [1, 0])
+ node_feature = graph.node_feat.get("node_feature")
+ triplet_indices = {
+ "idx_kj": graph.edge_feat["ti_idx_kj"].astype("int64"),
+ "idx_ji": graph.edge_feat["ti_idx_ji"].astype("int64"),
+ "idx_lk": graph.edge_feat["ti_idx_lk"].astype("int64"),
+ "idx_triplet": graph.edge_feat["ti_idx_triplet"].astype("int64"),
+ }
+
+ if self.use_extra_node_feature and node_feature is not None:
+ extra_node_feature = self.extra_emb(node_feature)
+ else:
+ extra_node_feature = None
+
+ num_nodes = z.shape[0]
+ dist, angle, torsion, i, j, idx_kj, idx_ji = compute_geometry(
+ pos, edge_index, triplet_indices
+ )
+
+ emb_out = self.emb_layer(dist, angle, torsion, idx_kj)
+
+ e = self.init_e(z, extra_node_feature, emb_out, i, j)
+ v = self.init_v(e, i, dim_size=num_nodes)
+ u = scatter_sum(v, node_batch, dim=0)
+
+ for update_e, update_v in zip(self.update_es, self.update_vs):
+ e = update_e(e, emb_out, idx_kj, idx_ji)
+ v = update_v(e, i, dim_size=num_nodes)
+ u = u + scatter_sum(v, node_batch, dim=0)
+
+ return u, pos
+
+ def forward(self, data, return_loss=True, return_prediction=True):
+ """Forward with the PaddleMaterials dict interface."""
+ assert (
+ return_loss or return_prediction
+ ), "At least one of return_loss or return_prediction must be True."
+
+ pred, pos = self._forward(data)
+
+ forces_pred = None
+ if self.energy_and_force:
+ grad = paddle.grad(pred.sum(), pos, create_graph=False, allow_unused=True)
+ if grad is not None and grad[0] is not None:
+ forces_pred = -grad[0]
+
+ loss_dict = {}
+ if return_loss:
+ label = data[self.property_name]
+ label_tensor = (
+ label.astype(paddle.get_default_dtype())
+ if isinstance(label, paddle.Tensor)
+ else paddle.to_tensor(label, dtype=paddle.get_default_dtype())
+ )
+ loss = paddle.nn.functional.l1_loss(pred, label_tensor)
+ loss_dict["loss"] = loss
+
+ if self.energy_and_force and forces_pred is not None:
+ force = data[self.force_key]
+ force_tensor = (
+ force.astype(paddle.get_default_dtype())
+ if isinstance(force, paddle.Tensor)
+ else paddle.to_tensor(force, dtype=paddle.get_default_dtype())
+ )
+ force_loss = paddle.nn.functional.l1_loss(forces_pred, force_tensor)
+ loss_dict["loss"] = loss + force_loss
+
+ prediction = {}
+ if return_prediction:
+ prediction[self.property_name] = pred
+ if self.energy_and_force:
+ if forces_pred is not None:
+ prediction[self.force_key] = forces_pred.detach()
+ else:
+ prediction[self.force_key] = paddle.zeros_like(pos)
+
+ return {"loss_dict": loss_dict, "pred_dict": prediction}
+
+ def predict(self, graphs):
+ """Inference interface for batch dicts or PGL graphs."""
+ if isinstance(graphs, list):
+ return [self.predict(graph) for graph in graphs]
+
+ data = graphs if isinstance(graphs, dict) else {"graph": graphs}
+ result = self.forward(data, return_loss=False, return_prediction=True)
+ return {
+ key: value.numpy() if isinstance(value, paddle.Tensor) else value
+ for key, value in result["pred_dict"].items()
+ }
diff --git a/ppmat/utils/scatter.py b/ppmat/utils/scatter.py
index b5de6495..db055487 100644
--- a/ppmat/utils/scatter.py
+++ b/ppmat/utils/scatter.py
@@ -31,6 +31,31 @@ def _broadcast(src: paddle.Tensor, other: paddle.Tensor, dim: int):
return src
+def scatter_argmin(
+ src: paddle.Tensor,
+ index: paddle.Tensor,
+ dim_size: Optional[int] = None,
+) -> paddle.Tensor:
+ """Return the source index of the minimum value in each group.
+
+ ``src`` and ``index`` must be one-dimensional. Empty groups are assigned
+ ``-1``. Ties are resolved by selecting the first occurrence in ``src``.
+ """
+ if src.ndim != 1 or index.ndim != 1 or src.shape[0] != index.shape[0]:
+ raise ValueError("src and index must be one-dimensional with equal length")
+
+ if dim_size is None:
+ dim_size = 0 if index.shape[0] == 0 else int(index.max()) + 1
+
+ out = paddle.full([dim_size], -1, dtype="int64")
+ if index.shape[0] == 0:
+ return out
+
+ order = paddle.argsort(src, stable=True)
+ groups, first = paddle.unique(index[order], return_index=True)
+ return paddle.scatter(out, groups, order[first], overwrite=True)
+
+
def _scatter_sum(
src: paddle.Tensor,
index: paddle.Tensor,
@@ -48,9 +73,18 @@ def _scatter_sum(
else:
size[dim] = int(index.max()) + 1
out = paddle.zeros(size, dtype=src.dtype)
- return paddle.put_along_axis(
- arr=out, indices=index, values=src, axis=dim, reduce="add"
- )
+ # FIXME: Paddle's put_along_axis backward (PutAlongAxisGradNode) crashes
+ # for dim=0; use one-hot + matmul as drop-in replacement.
+ if dim == 0:
+ # _broadcast expanded index to src.shape; collapse back to 1D via first column
+ idx_1d = index.reshape([-1, src.shape[1]])[:, 0] if index.ndim > 1 else index
+ one_hot = paddle.nn.functional.one_hot(idx_1d, out.shape[0]).cast(src.dtype)
+ # one_hot: [N, out_dim] -> [out_dim, N] @ [N, C] = [out_dim, C]
+ return paddle.mm(one_hot.t(), src)
+ else:
+ return paddle.put_along_axis(
+ arr=out, indices=index, values=src, axis=dim, reduce="add"
+ )
def _scatter_mean(
@@ -80,6 +114,28 @@ def _scatter_mean(
return out
+def _scatter_min(
+ src: paddle.Tensor,
+ index: paddle.Tensor,
+ dim: int = -1,
+ out: Optional[paddle.Tensor] = None,
+ dim_size: Optional[int] = None,
+) -> paddle.Tensor:
+ index = _broadcast(index, src, dim)
+ if out is None:
+ size = list(src.shape)
+ if dim_size is not None:
+ size[dim] = dim_size
+ elif index.numel() == 0:
+ size[dim] = 0
+ else:
+ size[dim] = int(index.max()) + 1
+ out = paddle.full(size, float("inf"), dtype=src.dtype)
+ return paddle.put_along_axis(
+ arr=out, indices=index, values=src, axis=dim, reduce="amin"
+ )
+
+
def scatter(
src: paddle.Tensor,
index: paddle.Tensor,
@@ -95,8 +151,10 @@ def scatter(
return _scatter_sum(src, index, dim, out, dim_size)
elif reduce == "mean":
return _scatter_mean(src, index, dim, out, dim_size)
+ elif reduce == "min":
+ return _scatter_min(src, index, dim, out, dim_size)
else:
- raise ValueError("Only support add or mean")
+ raise ValueError("Only support add, mean, or min")
def scatter_mean(
diff --git a/ppmat/utils/tests/test_scatter.py b/ppmat/utils/tests/test_scatter.py
new file mode 100644
index 00000000..0d4b92ac
--- /dev/null
+++ b/ppmat/utils/tests/test_scatter.py
@@ -0,0 +1,45 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import paddle
+
+from ppmat.utils.scatter import scatter_argmin
+
+
+def test_scatter_argmin_handles_unsorted_and_empty_groups():
+ values = paddle.to_tensor([3.0, -2.0, 4.0, -5.0, 1.0])
+ groups = paddle.to_tensor([2, 0, 2, 0, 2], dtype="int64")
+
+ result = scatter_argmin(values, groups, dim_size=4)
+
+ np.testing.assert_array_equal(result.numpy(), [3, -1, 4, -1])
+
+
+def test_scatter_argmin_selects_first_value_on_ties():
+ values = paddle.to_tensor([2.0, 1.0, 1.0, 3.0])
+ groups = paddle.to_tensor([0, 0, 0, 1], dtype="int64")
+
+ result = scatter_argmin(values, groups)
+
+ np.testing.assert_array_equal(result.numpy(), [1, 3])
+
+
+def test_scatter_argmin_handles_empty_input():
+ values = paddle.empty([0], dtype="float32")
+ groups = paddle.empty([0], dtype="int64")
+
+ result = scatter_argmin(values, groups, dim_size=3)
+
+ np.testing.assert_array_equal(result.numpy(), [-1, -1, -1])
diff --git a/property_prediction/README.md b/property_prediction/README.md
index 4d3e265e..5bef54bb 100644
--- a/property_prediction/README.md
+++ b/property_prediction/README.md
@@ -6,39 +6,53 @@ Property Prediction (PP) targets rapid, first-principles-level estimation of key
## 2.Models Matrix
-| **Supported Functions** | **[MEGNet](./configs/megnet/README.md)** | **[Comfomer](./configs/comformer/README.md)** | **GemNet** | **[DimeNet++](./configs/dimenet++/README.md)** |
-| -------------------------------------------- | :--------------------------------------: | :-------------------------------------------: | :--------: | :--------------------------------------------: |
-| **Forward Prediction · Materials Properties**| | | | |
-| Formation energy | ✅ | ✅ | 🚧 | ✅ |
-| Band gap | ✅ | ✅ | 🚧 | ✅ |
-| Bulk modulus | ✅ | ✅ | 🚧 | ✅ |
-| Shear modulus | ✅ | ✅ | 🚧 | ✅ |
-| Young’s modulus | ✅ | ✅ | 🚧 | ✅ |
-| Adsorption energy | 🚧 | 🚧 | 🚧 | 🚧 |
-| 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** | | | | |
-| MP2024 | ✅ | ✅ | — | — |
-| MP2020 | ✅ | ✅ | — | — |
-| MP2018 | ✅ | ✅ | 🚧 | — |
-| **JARVIS** | | | | |
-| dft_2d | ✅ | ✅ | — | ✅ |
-| dft_3d | ✅ | ✅ | — | — |
-| **Alexandria** | | | | |
-| pbe_2d | ✅ | ✅ | 🚧 | — |
-| **ML2DDB🌟** | ✅ | ✅ | ✅ | ✅ |
+| **Supported Functions** | **[MEGNet](./configs/megnet/README.md)** | **[Comfomer](./configs/comformer/README.md)** | **[DimeNet++](./configs/dimenet++/README.md)** | **[SphereNet](./configs/spherenet/README.md)** |
+| -------------------------------------------- | :--------------------------------------: | :-------------------------------------------: | :--------------------------------------------: | :--------------------------------------------: |
+| **Forward Prediction · Materials Properties**| | | | |
+| Formation energy | ✅ | ✅ | ✅ | — |
+| Band gap | ✅ | ✅ | ✅ | — |
+| Bulk modulus | ✅ | ✅ | ✅ | — |
+| Shear modulus | ✅ | ✅ | ✅ | — |
+| Young’s modulus | ✅ | ✅ | ✅ | — |
+| Adsorption energy | 🚧 | 🚧 | 🚧 | — |
+| **Forward Prediction · Molecular Properties**| | | | |
+| $\mu$ (dipole moment) | — | — | — | ✅ |
+| $\alpha$ (isotropic polarizability) | — | — | — | ✅ |
+| $\varepsilon_{\text{HOMO}}$ | — | — | — | ✅ |
+| $\varepsilon_{\text{LUMO}}$ | — | — | — | ✅ |
+| $\Delta\varepsilon$ (HOMO-LUMO gap) | — | — | — | ✅ |
+| $\langle R^2 \rangle$ (electronic spatial extent) | — | — | — | ✅ |
+| ZPVE (zero-point vibrational energy) | — | — | — | ✅ |
+| $U_0$ (internal energy at 0 K) | — | — | — | ✅ |
+| $U$ (internal energy at 298.15 K) | — | — | — | ✅ |
+| $H$ (enthalpy at 298.15 K) | — | — | — | ✅ |
+| $G$ (free energy at 298.15 K) | — | — | — | ✅ |
+| $C_v$ (heat capacity) | — | — | — | ✅ |
+| **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** | | | | |
+| MP2024 | ✅ | ✅ | — | — |
+| MP2020 | ✅ | ✅ | — | — |
+| MP2018 | ✅ | ✅ | — | — |
+| **JARVIS** | | | | |
+| dft_2d | ✅ | ✅ | ✅ | — |
+| dft_3d | ✅ | ✅ | — | — |
+| **Alexandria** | | | | |
+| pbe_2d | ✅ | ✅ | — | — |
+| **ML2DDB🌟** | ✅ | ✅ | ✅ | — |
+| **QM9** | — | — | ✅ | ✅ |
+| **MD17** | — | — | — | ✅ |
**Notice**:🌟 represent originate research work published from paddlematerials toolkit
diff --git a/property_prediction/__init__.py b/property_prediction/__init__.py
new file mode 100644
index 00000000..c0bee9ea
--- /dev/null
+++ b/property_prediction/__init__.py
@@ -0,0 +1 @@
+"""Public property-prediction interfaces."""
diff --git a/property_prediction/configs/gmtnet/README.md b/property_prediction/configs/gmtnet/README.md
new file mode 100644
index 00000000..a194cbec
--- /dev/null
+++ b/property_prediction/configs/gmtnet/README.md
@@ -0,0 +1,79 @@
+# GMTNet dielectric prediction
+
+This configuration integrates GMTNet into the PaddleMaterials public
+property-prediction path. GMTNet predicts a `3 x 3` dielectric tensor from a
+crystal structure.
+
+## Data and training
+
+Verify or reproduce the fixed-seed split for the normalized JARVIS dielectric
+dataset with the repository tool:
+
+```bash
+python -m ppmat.datasets.split_gmtnet_dataset --help
+```
+
+The canonical split is
+`property_prediction/configs/gmtnet/split_gmtnet_dielectric_seed32.json` and
+is installed as a `property_prediction` package resource. Use the public
+trainer entry point and this configuration for training or evaluation; the
+former standalone synthetic `train.py` smoke script is not a public GMTNet
+workflow. The YAML relies on `GMTNetDielectricDataset`'s resource fallback;
+callers can still pass an explicit `split_path`. Checkpoints are deliberately
+not committed to Git.
+
+```bash
+python property_prediction/train.py \
+ -c property_prediction/configs/gmtnet/gmtnet_jarvis_dielectric.yaml
+```
+
+GMTNet's inference path requires forward gradients internally. Keep
+`Predict.eval_with_no_grad: false` for this model, even during CPU evaluation.
+
+## Converted checkpoint prediction
+
+```python
+from property_prediction.predict import PropertyPredictor
+
+predictor = PropertyPredictor(
+ config_path="property_prediction/configs/gmtnet/gmtnet_jarvis_dielectric.yaml",
+ checkpoint_path="/path/to/gmtnet_checkpoint.pdparams",
+)
+```
+
+`predictor.from_structures(structure)` accepts one pymatgen `Structure`.
+An ordered non-empty list of structures is also supported. Low-level GMTNet
+Mapping inputs contain `graph`, `feature_mask`, and `matrix_equal` and can be
+passed to `predictor.model.predict`.
+
+## CIF contract
+
+GMTNet CIF prediction uses a model-specific precision-preserving reader:
+
+```python
+Structure.from_file(
+ path,
+ primitive=False,
+ sort=False,
+ merge_tol=0.0,
+ frac_tolerance=0.0,
+)
+```
+
+No primitive/conventional-cell conversion, site sort, site merge, or parser
+coordinate idealization is applied. The output tensor is expressed in the
+Cartesian frame of the parsed CIF structure. A CIF cannot recover floating
+point coordinates that were not encoded in its text.
+
+For a CIF directory, all files are parsed before any prediction begins. A
+malformed member therefore raises without producing partial directory
+predictions. Existing directory enumeration order is retained and is not a
+lexicographic-order guarantee.
+
+## Known numerical limits
+
+The float32 edge-order candidate and rotation canonicalization candidate were
+audited but are not applied. Equivalent edge-array orders can produce small
+float32 prediction differences, and symmetry-derived constraints remain
+sensitive to numerically equivalent rotation representations. These limits do
+not alter the precision-preserving CIF parsing contract.
diff --git a/property_prediction/configs/gmtnet/gmtnet_jarvis_dielectric.yaml b/property_prediction/configs/gmtnet/gmtnet_jarvis_dielectric.yaml
new file mode 100644
index 00000000..c2d3e9c6
--- /dev/null
+++ b/property_prediction/configs/gmtnet/gmtnet_jarvis_dielectric.yaml
@@ -0,0 +1,145 @@
+Global:
+ label_names: ["dielectric"]
+ do_train: true
+ do_eval: true
+ do_test: true
+
+Trainer:
+ max_epochs: 10
+ seed: 32
+ output_dir: ./output/gmtnet_dielectric
+ save_freq: 10
+ log_freq: 10
+ start_eval_epoch: 1
+ eval_freq: 1
+ pretrained_model_path: null
+ pretrained_weight_name: null
+ resume_from_checkpoint: null
+ use_amp: false
+ amp_level: "O1"
+ eval_with_no_grad: false
+ gradient_accumulation_steps: 1
+ best_metric_indicator: "eval_metric"
+ name_for_best_metric: "dielectric"
+ 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__: GMTNet
+ __init_params__:
+ args:
+ atom_input_features: 92
+ edge_features: 512
+ embedding_features: 128
+ output_features: 9
+ num_layers: 2
+ target: dielectric
+ use_mask: false
+ reduce_cell: false
+ loss_cfg:
+ __class_name__: MSELoss
+ __init_params__:
+ reduction: mean
+
+Optimizer:
+ __class_name__: AdamW
+ __init_params__:
+ lr: 1.0e-5
+ weight_decay: 1.0e-5
+
+Metric:
+ dielectric:
+ __class_name__: paddle.nn.L1Loss
+ __init_params__: {}
+
+Dataset:
+ train:
+ dataset:
+ __class_name__: GMTNetDielectricDataset
+ __init_params__:
+ data_path: data/gmtnet/gmtnet_dielectric_normalized_v1.pkl
+ split: train
+ verify_sha256: true
+ build_graph_cfg:
+ __class_name__: GMTNetGraphConverter
+ __init_params__:
+ cutoff: 4.0
+ max_neighbors: 16
+ atom_features: cgcnn
+ use_canonize: true
+ reduce_cell: false
+ loader:
+ num_workers: 0
+ use_shared_memory: false
+ collate_fn: DefaultCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 1
+ shuffle: true
+ drop_last: false
+ val:
+ dataset:
+ __class_name__: GMTNetDielectricDataset
+ __init_params__:
+ data_path: data/gmtnet/gmtnet_dielectric_normalized_v1.pkl
+ split: val
+ verify_sha256: true
+ build_graph_cfg:
+ __class_name__: GMTNetGraphConverter
+ __init_params__:
+ cutoff: 4.0
+ max_neighbors: 16
+ atom_features: cgcnn
+ use_canonize: true
+ reduce_cell: false
+ loader:
+ num_workers: 0
+ use_shared_memory: false
+ collate_fn: DefaultCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 1
+ shuffle: false
+ drop_last: false
+ test:
+ dataset:
+ __class_name__: GMTNetDielectricDataset
+ __init_params__:
+ data_path: data/gmtnet/gmtnet_dielectric_normalized_v1.pkl
+ split: test
+ verify_sha256: true
+ build_graph_cfg:
+ __class_name__: GMTNetGraphConverter
+ __init_params__:
+ cutoff: 4.0
+ max_neighbors: 16
+ atom_features: cgcnn
+ use_canonize: true
+ reduce_cell: false
+ loader:
+ num_workers: 0
+ use_shared_memory: false
+ collate_fn: DefaultCollator
+ sampler:
+ __class_name__: BatchSampler
+ __init_params__:
+ batch_size: 1
+ shuffle: false
+ drop_last: false
+
+Predict:
+ graph_converter:
+ __class_name__: GMTNetGraphConverter
+ __init_params__:
+ cutoff: 4.0
+ max_neighbors: 16
+ atom_features: cgcnn
+ use_canonize: true
+ reduce_cell: false
+ eval_with_no_grad: false
diff --git a/property_prediction/configs/gmtnet/split_gmtnet_dielectric_seed32.json b/property_prediction/configs/gmtnet/split_gmtnet_dielectric_seed32.json
new file mode 100644
index 00000000..c67400f5
--- /dev/null
+++ b/property_prediction/configs/gmtnet/split_gmtnet_dielectric_seed32.json
@@ -0,0 +1,4739 @@
+{
+ "generation_method": "torch.utils.data.random_split",
+ "normalized_dataset_sha256": "eb0b9516c937575afe3f20a0f88953724abfcde6de07c4af15248468598b349f",
+ "normalized_schema_version": 1,
+ "num_records": 4713,
+ "schema_version": 1,
+ "seed": 32,
+ "source_original_dataset_sha256": "5a2198f51f4a7f9aa26fa6be60ed65db0ecc0a13646d63399e897b8168dcbb4d",
+ "split_indices_sha256": {
+ "test": "23bf024dede3d367130dbf74fb0d4f5e70e4aaaacf68172ceca1d7586b903470",
+ "train": "6fc27a45d6c911bc4df0b6687bdb0e38fe825d0aed7b8dc4c912afb458857010",
+ "val": "3907ab990e5202ede42d044bf6eab47b2c7f1280688f3d3ae082b416fc7acb11"
+ },
+ "split_sizes": {
+ "test": 472,
+ "train": 3770,
+ "val": 471
+ },
+ "test_indices": [
+ 747,
+ 1423,
+ 1322,
+ 856,
+ 2521,
+ 2188,
+ 2308,
+ 1878,
+ 3544,
+ 653,
+ 1098,
+ 2646,
+ 178,
+ 3712,
+ 1179,
+ 59,
+ 3714,
+ 3559,
+ 1378,
+ 1494,
+ 3844,
+ 4596,
+ 4476,
+ 2955,
+ 1095,
+ 1340,
+ 3425,
+ 475,
+ 2949,
+ 3471,
+ 2502,
+ 120,
+ 966,
+ 1320,
+ 3403,
+ 3625,
+ 501,
+ 185,
+ 4304,
+ 247,
+ 283,
+ 3091,
+ 3348,
+ 710,
+ 556,
+ 1188,
+ 3331,
+ 646,
+ 801,
+ 2979,
+ 2432,
+ 3701,
+ 3852,
+ 3684,
+ 3334,
+ 988,
+ 1455,
+ 156,
+ 1590,
+ 1005,
+ 3375,
+ 3459,
+ 3927,
+ 4710,
+ 3096,
+ 2583,
+ 4302,
+ 3101,
+ 3823,
+ 3258,
+ 1897,
+ 1499,
+ 269,
+ 1717,
+ 146,
+ 1457,
+ 1913,
+ 4677,
+ 3915,
+ 2246,
+ 1584,
+ 3579,
+ 3965,
+ 281,
+ 1651,
+ 2928,
+ 4120,
+ 2957,
+ 1012,
+ 1099,
+ 3434,
+ 2057,
+ 3063,
+ 667,
+ 1567,
+ 2100,
+ 2916,
+ 4042,
+ 3156,
+ 734,
+ 754,
+ 623,
+ 96,
+ 2985,
+ 571,
+ 1357,
+ 384,
+ 1183,
+ 1995,
+ 2317,
+ 2433,
+ 2913,
+ 4033,
+ 4693,
+ 3338,
+ 301,
+ 491,
+ 1483,
+ 2679,
+ 3937,
+ 4196,
+ 4568,
+ 4188,
+ 1863,
+ 2343,
+ 4634,
+ 4419,
+ 2486,
+ 3468,
+ 1276,
+ 3699,
+ 4254,
+ 2518,
+ 1555,
+ 2818,
+ 762,
+ 3185,
+ 2755,
+ 335,
+ 4122,
+ 601,
+ 4171,
+ 4390,
+ 4108,
+ 4444,
+ 22,
+ 89,
+ 497,
+ 2167,
+ 3969,
+ 3330,
+ 2726,
+ 3647,
+ 1245,
+ 4541,
+ 3076,
+ 392,
+ 4704,
+ 2754,
+ 815,
+ 2161,
+ 1168,
+ 1380,
+ 2825,
+ 640,
+ 2748,
+ 1094,
+ 196,
+ 1079,
+ 2908,
+ 3443,
+ 170,
+ 4566,
+ 777,
+ 943,
+ 954,
+ 2187,
+ 3437,
+ 2284,
+ 3780,
+ 2402,
+ 1299,
+ 663,
+ 1629,
+ 1323,
+ 3901,
+ 969,
+ 1923,
+ 2254,
+ 4208,
+ 3382,
+ 1238,
+ 1283,
+ 4089,
+ 1065,
+ 3648,
+ 1953,
+ 2736,
+ 3159,
+ 2903,
+ 1634,
+ 3855,
+ 4386,
+ 2963,
+ 4610,
+ 1614,
+ 2654,
+ 1068,
+ 2385,
+ 375,
+ 4015,
+ 2013,
+ 4412,
+ 2808,
+ 1026,
+ 3030,
+ 4308,
+ 191,
+ 236,
+ 1515,
+ 2633,
+ 3896,
+ 1613,
+ 760,
+ 2466,
+ 4101,
+ 2429,
+ 396,
+ 530,
+ 4350,
+ 336,
+ 3929,
+ 4462,
+ 2554,
+ 3737,
+ 3970,
+ 4531,
+ 85,
+ 2665,
+ 3781,
+ 3871,
+ 496,
+ 362,
+ 583,
+ 3821,
+ 1550,
+ 3591,
+ 223,
+ 4315,
+ 1571,
+ 3733,
+ 2636,
+ 855,
+ 2552,
+ 3848,
+ 280,
+ 2638,
+ 722,
+ 4446,
+ 2769,
+ 3390,
+ 4130,
+ 458,
+ 2860,
+ 3082,
+ 2293,
+ 1083,
+ 4373,
+ 2288,
+ 91,
+ 1233,
+ 235,
+ 346,
+ 3804,
+ 1663,
+ 665,
+ 244,
+ 3151,
+ 112,
+ 2223,
+ 87,
+ 3492,
+ 133,
+ 4555,
+ 4261,
+ 1165,
+ 3458,
+ 2471,
+ 1730,
+ 681,
+ 1531,
+ 3635,
+ 2040,
+ 1794,
+ 3668,
+ 2833,
+ 3809,
+ 961,
+ 24,
+ 1197,
+ 542,
+ 382,
+ 3549,
+ 4367,
+ 1566,
+ 3343,
+ 2355,
+ 1067,
+ 3790,
+ 1089,
+ 317,
+ 2264,
+ 67,
+ 3496,
+ 2899,
+ 763,
+ 1387,
+ 4553,
+ 1223,
+ 3046,
+ 1498,
+ 4563,
+ 4247,
+ 2737,
+ 4493,
+ 1981,
+ 332,
+ 3165,
+ 1689,
+ 605,
+ 1404,
+ 4025,
+ 271,
+ 3876,
+ 506,
+ 1767,
+ 2711,
+ 3097,
+ 3488,
+ 2143,
+ 1756,
+ 4570,
+ 4486,
+ 714,
+ 441,
+ 2992,
+ 4323,
+ 3181,
+ 2062,
+ 1670,
+ 2468,
+ 3463,
+ 2109,
+ 1894,
+ 1101,
+ 635,
+ 189,
+ 3012,
+ 2602,
+ 1226,
+ 2251,
+ 4552,
+ 2587,
+ 474,
+ 2259,
+ 3016,
+ 4509,
+ 357,
+ 807,
+ 4543,
+ 4193,
+ 3910,
+ 897,
+ 3213,
+ 4626,
+ 3585,
+ 507,
+ 852,
+ 3342,
+ 842,
+ 4206,
+ 1540,
+ 3995,
+ 1843,
+ 4005,
+ 1022,
+ 4199,
+ 894,
+ 4522,
+ 951,
+ 4646,
+ 502,
+ 827,
+ 2031,
+ 3588,
+ 1800,
+ 3044,
+ 3237,
+ 1935,
+ 2830,
+ 1845,
+ 4569,
+ 3278,
+ 720,
+ 3374,
+ 697,
+ 3248,
+ 2898,
+ 2487,
+ 155,
+ 459,
+ 4177,
+ 3750,
+ 4702,
+ 293,
+ 2612,
+ 1078,
+ 2048,
+ 214,
+ 4179,
+ 4104,
+ 536,
+ 719,
+ 2533,
+ 3064,
+ 2410,
+ 3084,
+ 4292,
+ 2342,
+ 2353,
+ 3601,
+ 1112,
+ 2785,
+ 832,
+ 208,
+ 162,
+ 310,
+ 3863,
+ 4106,
+ 4070,
+ 533,
+ 3893,
+ 2083,
+ 4542,
+ 1699,
+ 4429,
+ 203,
+ 3771,
+ 2799,
+ 1548,
+ 4141,
+ 1868,
+ 1914,
+ 493,
+ 3355,
+ 1208,
+ 2959,
+ 795,
+ 1057,
+ 2765,
+ 3045,
+ 1196,
+ 1150,
+ 2988,
+ 3866,
+ 1536,
+ 2140,
+ 75,
+ 4249,
+ 2089,
+ 1696,
+ 1558
+ ],
+ "torch_version": "2.6.0+cu124",
+ "train_indices": [
+ 3244,
+ 3052,
+ 3259,
+ 1107,
+ 3797,
+ 567,
+ 216,
+ 258,
+ 494,
+ 4272,
+ 535,
+ 1570,
+ 2616,
+ 1782,
+ 766,
+ 3038,
+ 1775,
+ 3252,
+ 3327,
+ 2572,
+ 2570,
+ 2131,
+ 1237,
+ 2049,
+ 767,
+ 4075,
+ 2067,
+ 3581,
+ 3155,
+ 3961,
+ 4580,
+ 4088,
+ 2331,
+ 2706,
+ 2306,
+ 70,
+ 29,
+ 1707,
+ 4093,
+ 4515,
+ 1824,
+ 35,
+ 1684,
+ 1866,
+ 2730,
+ 3264,
+ 3164,
+ 33,
+ 4203,
+ 2871,
+ 3184,
+ 2493,
+ 2174,
+ 1274,
+ 3103,
+ 2677,
+ 1202,
+ 1019,
+ 3203,
+ 409,
+ 3373,
+ 2126,
+ 1481,
+ 1108,
+ 3484,
+ 3224,
+ 3795,
+ 3550,
+ 3095,
+ 2828,
+ 3765,
+ 1906,
+ 4276,
+ 875,
+ 2363,
+ 3174,
+ 2035,
+ 2921,
+ 3347,
+ 2055,
+ 1650,
+ 755,
+ 1405,
+ 3862,
+ 1232,
+ 211,
+ 2781,
+ 2365,
+ 3047,
+ 3570,
+ 328,
+ 2329,
+ 2943,
+ 1719,
+ 1048,
+ 3553,
+ 1049,
+ 1400,
+ 3690,
+ 3691,
+ 2930,
+ 2891,
+ 1427,
+ 4348,
+ 4572,
+ 2104,
+ 1062,
+ 3711,
+ 49,
+ 1192,
+ 153,
+ 1597,
+ 1572,
+ 352,
+ 3154,
+ 1045,
+ 3742,
+ 2403,
+ 403,
+ 1187,
+ 2856,
+ 1937,
+ 2382,
+ 4538,
+ 2483,
+ 2253,
+ 4707,
+ 69,
+ 360,
+ 1146,
+ 804,
+ 2610,
+ 2125,
+ 197,
+ 4651,
+ 2549,
+ 46,
+ 4336,
+ 1395,
+ 1587,
+ 1778,
+ 3043,
+ 1811,
+ 1109,
+ 3938,
+ 1313,
+ 4098,
+ 3722,
+ 2310,
+ 2700,
+ 4147,
+ 1159,
+ 31,
+ 2178,
+ 2966,
+ 92,
+ 4482,
+ 521,
+ 3739,
+ 3436,
+ 1044,
+ 2206,
+ 3730,
+ 3541,
+ 585,
+ 2299,
+ 882,
+ 2581,
+ 4151,
+ 4266,
+ 4473,
+ 3515,
+ 1672,
+ 510,
+ 1033,
+ 2300,
+ 2862,
+ 4103,
+ 2182,
+ 3048,
+ 1338,
+ 1267,
+ 1724,
+ 3913,
+ 79,
+ 1462,
+ 848,
+ 4311,
+ 4125,
+ 364,
+ 1693,
+ 3759,
+ 2298,
+ 4394,
+ 837,
+ 1627,
+ 3612,
+ 2227,
+ 2366,
+ 2164,
+ 3104,
+ 204,
+ 1151,
+ 1871,
+ 2747,
+ 2811,
+ 3592,
+ 407,
+ 2337,
+ 1559,
+ 2990,
+ 3251,
+ 3324,
+ 1591,
+ 11,
+ 2105,
+ 4405,
+ 3394,
+ 4212,
+ 262,
+ 4558,
+ 4061,
+ 3299,
+ 4485,
+ 3836,
+ 3903,
+ 1161,
+ 4331,
+ 2361,
+ 2387,
+ 2912,
+ 3410,
+ 972,
+ 973,
+ 2504,
+ 1968,
+ 677,
+ 4508,
+ 2271,
+ 718,
+ 2037,
+ 2460,
+ 3476,
+ 2794,
+ 3417,
+ 787,
+ 4230,
+ 1001,
+ 119,
+ 560,
+ 401,
+ 620,
+ 1411,
+ 4099,
+ 1989,
+ 1793,
+ 3968,
+ 520,
+ 4584,
+ 3829,
+ 4094,
+ 3603,
+ 4111,
+ 2200,
+ 184,
+ 1886,
+ 485,
+ 3747,
+ 4502,
+ 694,
+ 1785,
+ 3041,
+ 1787,
+ 4401,
+ 761,
+ 477,
+ 3485,
+ 2279,
+ 461,
+ 1174,
+ 4687,
+ 999,
+ 1628,
+ 4271,
+ 3227,
+ 4384,
+ 831,
+ 4062,
+ 1144,
+ 1851,
+ 975,
+ 3802,
+ 377,
+ 32,
+ 2670,
+ 3982,
+ 1796,
+ 3306,
+ 1105,
+ 2144,
+ 3015,
+ 1618,
+ 3596,
+ 2738,
+ 776,
+ 2129,
+ 3514,
+ 3607,
+ 4045,
+ 4190,
+ 2698,
+ 228,
+ 237,
+ 4215,
+ 1449,
+ 3646,
+ 3746,
+ 3629,
+ 4016,
+ 2377,
+ 95,
+ 4643,
+ 2275,
+ 1988,
+ 2154,
+ 2065,
+ 711,
+ 1519,
+ 4004,
+ 470,
+ 996,
+ 492,
+ 3621,
+ 2553,
+ 3346,
+ 132,
+ 312,
+ 3675,
+ 1885,
+ 4270,
+ 3308,
+ 3560,
+ 633,
+ 2540,
+ 4607,
+ 4066,
+ 2822,
+ 1042,
+ 54,
+ 4250,
+ 547,
+ 2688,
+ 445,
+ 3397,
+ 909,
+ 222,
+ 1593,
+ 2448,
+ 1394,
+ 4507,
+ 1173,
+ 3637,
+ 599,
+ 115,
+ 2805,
+ 1733,
+ 4447,
+ 2690,
+ 939,
+ 4174,
+ 2563,
+ 1945,
+ 574,
+ 4518,
+ 2147,
+ 1304,
+ 200,
+ 1899,
+ 775,
+ 4418,
+ 576,
+ 2074,
+ 4213,
+ 3627,
+ 3148,
+ 4055,
+ 3088,
+ 2106,
+ 537,
+ 2536,
+ 2831,
+ 4423,
+ 201,
+ 463,
+ 611,
+ 2199,
+ 652,
+ 0,
+ 657,
+ 1252,
+ 1358,
+ 4277,
+ 2053,
+ 1617,
+ 435,
+ 2027,
+ 2629,
+ 2376,
+ 2579,
+ 4624,
+ 254,
+ 1491,
+ 4635,
+ 3239,
+ 4698,
+ 291,
+ 3740,
+ 896,
+ 1305,
+ 4267,
+ 1757,
+ 1364,
+ 159,
+ 3431,
+ 1944,
+ 1747,
+ 3037,
+ 1640,
+ 198,
+ 1537,
+ 1495,
+ 2444,
+ 3123,
+ 3872,
+ 2364,
+ 2606,
+ 1341,
+ 1125,
+ 2177,
+ 2590,
+ 4370,
+ 3911,
+ 3897,
+ 3365,
+ 1583,
+ 3539,
+ 2139,
+ 1947,
+ 3816,
+ 2237,
+ 3479,
+ 221,
+ 3408,
+ 255,
+ 2756,
+ 3460,
+ 658,
+ 3775,
+ 2595,
+ 499,
+ 4191,
+ 582,
+ 3208,
+ 3709,
+ 4533,
+ 3632,
+ 3812,
+ 3870,
+ 1119,
+ 1329,
+ 3412,
+ 2971,
+ 3315,
+ 277,
+ 2938,
+ 1435,
+ 4131,
+ 1644,
+ 139,
+ 4692,
+ 1193,
+ 1568,
+ 466,
+ 4484,
+ 4256,
+ 4638,
+ 250,
+ 270,
+ 1542,
+ 18,
+ 3432,
+ 3391,
+ 3332,
+ 619,
+ 47,
+ 215,
+ 3292,
+ 2489,
+ 1085,
+ 581,
+ 373,
+ 1849,
+ 4468,
+ 1718,
+ 1436,
+ 3068,
+ 4586,
+ 3760,
+ 420,
+ 3414,
+ 4069,
+ 812,
+ 1416,
+ 569,
+ 2668,
+ 1138,
+ 2152,
+ 678,
+ 702,
+ 1728,
+ 942,
+ 1339,
+ 116,
+ 1284,
+ 4636,
+ 630,
+ 1683,
+ 3212,
+ 2705,
+ 858,
+ 4349,
+ 1203,
+ 3891,
+ 4291,
+ 1698,
+ 1791,
+ 2479,
+ 614,
+ 985,
+ 1023,
+ 2881,
+ 511,
+ 3207,
+ 3761,
+ 3880,
+ 615,
+ 3428,
+ 586,
+ 261,
+ 390,
+ 1807,
+ 4690,
+ 464,
+ 3034,
+ 1227,
+ 88,
+ 2893,
+ 3444,
+ 3644,
+ 1859,
+ 2017,
+ 609,
+ 3236,
+ 1037,
+ 2001,
+ 2717,
+ 932,
+ 1077,
+ 597,
+ 1638,
+ 4225,
+ 219,
+ 3577,
+ 2485,
+ 2866,
+ 3439,
+ 2680,
+ 3513,
+ 983,
+ 997,
+ 735,
+ 4567,
+ 1772,
+ 1017,
+ 1501,
+ 3516,
+ 604,
+ 1122,
+ 1991,
+ 3813,
+ 325,
+ 1695,
+ 2556,
+ 2691,
+ 2102,
+ 273,
+ 2283,
+ 1870,
+ 1977,
+ 2853,
+ 4100,
+ 4616,
+ 3145,
+ 3105,
+ 3017,
+ 3462,
+ 1465,
+ 1502,
+ 3825,
+ 1243,
+ 4496,
+ 1965,
+ 4008,
+ 2802,
+ 518,
+ 330,
+ 3540,
+ 2134,
+ 2523,
+ 108,
+ 4594,
+ 3389,
+ 234,
+ 68,
+ 990,
+ 1588,
+ 1222,
+ 2559,
+ 3200,
+ 1948,
+ 2517,
+ 3361,
+ 260,
+ 911,
+ 4035,
+ 1980,
+ 2241,
+ 2791,
+ 3214,
+ 467,
+ 2132,
+ 647,
+ 1419,
+ 602,
+ 2598,
+ 4560,
+ 3307,
+ 1116,
+ 1539,
+ 3353,
+ 4660,
+ 4359,
+ 429,
+ 695,
+ 4170,
+ 2560,
+ 1928,
+ 3121,
+ 3815,
+ 1690,
+ 1549,
+ 3435,
+ 4289,
+ 3433,
+ 2920,
+ 3889,
+ 1189,
+ 3830,
+ 2401,
+ 127,
+ 3947,
+ 2133,
+ 1762,
+ 181,
+ 929,
+ 3614,
+ 1076,
+ 3572,
+ 3222,
+ 4112,
+ 3004,
+ 1241,
+ 398,
+ 1289,
+ 323,
+ 4671,
+ 4083,
+ 4209,
+ 449,
+ 737,
+ 3840,
+ 1319,
+ 4380,
+ 1582,
+ 3061,
+ 1534,
+ 3219,
+ 1626,
+ 2547,
+ 124,
+ 2322,
+ 1784,
+ 27,
+ 664,
+ 4435,
+ 3702,
+ 2773,
+ 2506,
+ 1912,
+ 3924,
+ 4381,
+ 4683,
+ 1382,
+ 1992,
+ 1596,
+ 1575,
+ 3153,
+ 1277,
+ 864,
+ 671,
+ 3135,
+ 3369,
+ 3026,
+ 2028,
+ 2180,
+ 337,
+ 4441,
+ 750,
+ 4158,
+ 912,
+ 431,
+ 2685,
+ 1734,
+ 316,
+ 4255,
+ 4195,
+ 233,
+ 410,
+ 1016,
+ 1661,
+ 4178,
+ 2185,
+ 1399,
+ 2934,
+ 488,
+ 659,
+ 4192,
+ 28,
+ 1035,
+ 989,
+ 4347,
+ 2166,
+ 2520,
+ 2669,
+ 3721,
+ 366,
+ 3053,
+ 4076,
+ 1009,
+ 2592,
+ 3728,
+ 329,
+ 4029,
+ 733,
+ 4402,
+ 1351,
+ 1755,
+ 4345,
+ 1268,
+ 1677,
+ 2260,
+ 2321,
+ 1836,
+ 1403,
+ 2775,
+ 3860,
+ 419,
+ 173,
+ 3328,
+ 4463,
+ 4239,
+ 2838,
+ 3851,
+ 1685,
+ 1974,
+ 769,
+ 3989,
+ 4024,
+ 3636,
+ 1604,
+ 1143,
+ 4332,
+ 618,
+ 2232,
+ 2019,
+ 1407,
+ 580,
+ 1036,
+ 333,
+ 2183,
+ 3505,
+ 4593,
+ 51,
+ 2362,
+ 3475,
+ 4454,
+ 541,
+ 1000,
+ 1008,
+ 2886,
+ 3615,
+ 498,
+ 4265,
+ 1440,
+ 3597,
+ 3705,
+ 2220,
+ 4411,
+ 517,
+ 2760,
+ 3616,
+ 171,
+ 2780,
+ 4202,
+ 3535,
+ 34,
+ 1397,
+ 4356,
+ 324,
+ 2936,
+ 3557,
+ 3192,
+ 2986,
+ 544,
+ 1854,
+ 4164,
+ 3257,
+ 118,
+ 3984,
+ 4358,
+ 2243,
+ 2474,
+ 44,
+ 4183,
+ 40,
+ 2931,
+ 2503,
+ 2076,
+ 1625,
+ 1654,
+ 2249,
+ 3085,
+ 2470,
+ 944,
+ 2196,
+ 2236,
+ 1027,
+ 1671,
+ 73,
+ 4595,
+ 3464,
+ 570,
+ 4133,
+ 950,
+ 4546,
+ 1113,
+ 1406,
+ 296,
+ 4395,
+ 1360,
+ 114,
+ 3215,
+ 2008,
+ 1229,
+ 2202,
+ 2155,
+ 2005,
+ 1508,
+ 515,
+ 736,
+ 2582,
+ 654,
+ 3383,
+ 715,
+ 287,
+ 1711,
+ 712,
+ 617,
+ 3831,
+ 164,
+ 1477,
+ 1561,
+ 205,
+ 3262,
+ 3149,
+ 1919,
+ 1492,
+ 3283,
+ 3998,
+ 3452,
+ 1983,
+ 427,
+ 2081,
+ 4673,
+ 2733,
+ 3639,
+ 100,
+ 4053,
+ 1393,
+ 3531,
+ 4556,
+ 607,
+ 1438,
+ 2660,
+ 3548,
+ 2640,
+ 853,
+ 4194,
+ 4397,
+ 3865,
+ 4023,
+ 2953,
+ 2820,
+ 1874,
+ 3217,
+ 1474,
+ 504,
+ 2868,
+ 2221,
+ 4128,
+ 1471,
+ 4365,
+ 1490,
+ 1951,
+ 3895,
+ 1181,
+ 1557,
+ 572,
+ 2851,
+ 1154,
+ 3057,
+ 4352,
+ 3359,
+ 4232,
+ 824,
+ 548,
+ 1213,
+ 4001,
+ 2235,
+ 1527,
+ 3357,
+ 2537,
+ 2935,
+ 1861,
+ 2550,
+ 4154,
+ 279,
+ 917,
+ 1327,
+ 1097,
+ 1987,
+ 331,
+ 1297,
+ 4238,
+ 4165,
+ 161,
+ 3653,
+ 3882,
+ 3792,
+ 1294,
+ 962,
+ 3379,
+ 765,
+ 3094,
+ 2295,
+ 370,
+ 3973,
+ 3778,
+ 141,
+ 2535,
+ 1310,
+ 1815,
+ 741,
+ 2645,
+ 1443,
+ 2009,
+ 4372,
+ 2023,
+ 4135,
+ 2635,
+ 967,
+ 4318,
+ 2225,
+ 380,
+ 4661,
+ 3842,
+ 257,
+ 4137,
+ 1832,
+ 2925,
+ 3589,
+ 854,
+ 4184,
+ 2163,
+ 1810,
+ 2311,
+ 993,
+ 451,
+ 2788,
+ 4393,
+ 2948,
+ 443,
+ 4124,
+ 4139,
+ 3912,
+ 1715,
+ 1758,
+ 3789,
+ 2087,
+ 1248,
+ 4450,
+ 2911,
+ 4353,
+ 4051,
+ 2900,
+ 4591,
+ 4487,
+ 4010,
+ 1955,
+ 1486,
+ 3900,
+ 4574,
+ 4523,
+ 1170,
+ 3243,
+ 2334,
+ 1674,
+ 2768,
+ 503,
+ 771,
+ 4059,
+ 1556,
+ 2804,
+ 2555,
+ 4328,
+ 3321,
+ 2002,
+ 2519,
+ 2724,
+ 2233,
+ 2653,
+ 4503,
+ 2369,
+ 4641,
+ 3051,
+ 1656,
+ 2354,
+ 3233,
+ 4161,
+ 2247,
+ 2103,
+ 2181,
+ 4228,
+ 252,
+ 1902,
+ 3011,
+ 1032,
+ 2422,
+ 1066,
+ 1736,
+ 4126,
+ 1839,
+ 1260,
+ 4153,
+ 905,
+ 36,
+ 1330,
+ 2324,
+ 3449,
+ 405,
+ 413,
+ 3785,
+ 2150,
+ 193,
+ 4065,
+ 3556,
+ 1994,
+ 3445,
+ 2855,
+ 522,
+ 4527,
+ 1727,
+ 925,
+ 1761,
+ 1774,
+ 2228,
+ 3916,
+ 2014,
+ 1745,
+ 538,
+ 2158,
+ 4408,
+ 1142,
+ 2142,
+ 1776,
+ 2061,
+ 3611,
+ 1949,
+ 1489,
+ 3054,
+ 3282,
+ 2884,
+ 3826,
+ 2043,
+ 1288,
+ 1090,
+ 1599,
+ 2240,
+ 870,
+ 3628,
+ 2749,
+ 145,
+ 2524,
+ 3499,
+ 3036,
+ 2390,
+ 593,
+ 4002,
+ 3738,
+ 160,
+ 3936,
+ 3232,
+ 436,
+ 1503,
+ 2291,
+ 898,
+ 1190,
+ 1891,
+ 3534,
+ 306,
+ 2124,
+ 2850,
+ 2901,
+ 1530,
+ 836,
+ 1362,
+ 4612,
+ 4399,
+ 2514,
+ 1847,
+ 2408,
+ 4320,
+ 314,
+ 1328,
+ 3083,
+ 2406,
+ 4466,
+ 1392,
+ 2020,
+ 209,
+ 2110,
+ 4364,
+ 2615,
+ 3079,
+ 1386,
+ 1780,
+ 4529,
+ 3060,
+ 2418,
+ 411,
+ 4459,
+ 3430,
+ 4424,
+ 3441,
+ 976,
+ 3281,
+ 3843,
+ 2643,
+ 2952,
+ 2883,
+ 1153,
+ 460,
+ 1647,
+ 843,
+ 3808,
+ 2735,
+ 1282,
+ 1295,
+ 3708,
+ 1704,
+ 276,
+ 2025,
+ 2923,
+ 417,
+ 3421,
+ 2168,
+ 2915,
+ 2507,
+ 2060,
+ 2459,
+ 2367,
+ 20,
+ 4391,
+ 1266,
+ 886,
+ 1074,
+ 3301,
+ 10,
+ 2962,
+ 3137,
+ 4539,
+ 4241,
+ 3482,
+ 4605,
+ 2098,
+ 1817,
+ 224,
+ 2409,
+ 3763,
+ 4245,
+ 742,
+ 2975,
+ 4175,
+ 525,
+ 1869,
+ 3692,
+ 1603,
+ 2453,
+ 797,
+ 3558,
+ 4341,
+ 3533,
+ 2258,
+ 4007,
+ 4577,
+ 292,
+ 687,
+ 131,
+ 3873,
+ 2285,
+ 21,
+ 2244,
+ 692,
+ 4519,
+ 4672,
+ 4342,
+ 2046,
+ 4309,
+ 1538,
+ 309,
+ 1391,
+ 62,
+ 4689,
+ 2499,
+ 4420,
+ 1039,
+ 4259,
+ 2270,
+ 991,
+ 2075,
+ 4143,
+ 4145,
+ 4430,
+ 56,
+ 3163,
+ 344,
+ 2647,
+ 3524,
+ 2836,
+ 649,
+ 150,
+ 971,
+ 1493,
+ 1333,
+ 4160,
+ 845,
+ 1086,
+ 584,
+ 2117,
+ 1544,
+ 481,
+ 1667,
+ 4050,
+ 2007,
+ 53,
+ 4483,
+ 3167,
+ 408,
+ 2865,
+ 1217,
+ 1141,
+ 4579,
+ 4166,
+ 1061,
+ 4187,
+ 2758,
+ 3112,
+ 1255,
+ 4064,
+ 977,
+ 4231,
+ 1770,
+ 423,
+ 709,
+ 3351,
+ 2392,
+ 379,
+ 1769,
+ 3674,
+ 3892,
+ 1242,
+ 4324,
+ 149,
+ 2018,
+ 2093,
+ 1971,
+ 1660,
+ 1047,
+ 4583,
+ 2335,
+ 45,
+ 4058,
+ 3413,
+ 1199,
+ 2348,
+ 2442,
+ 3850,
+ 3368,
+ 288,
+ 4319,
+ 1560,
+ 210,
+ 786,
+ 1653,
+ 30,
+ 3861,
+ 2630,
+ 3877,
+ 3925,
+ 953,
+ 1850,
+ 3593,
+ 4078,
+ 2250,
+ 2307,
+ 3172,
+ 4398,
+ 3206,
+ 248,
+ 3966,
+ 2391,
+ 1450,
+ 2929,
+ 3009,
+ 1998,
+ 1901,
+ 2932,
+ 3344,
+ 610,
+ 1754,
+ 3396,
+ 4301,
+ 2739,
+ 4662,
+ 561,
+ 4481,
+ 3729,
+ 2287,
+ 591,
+ 142,
+ 2672,
+ 1898,
+ 3972,
+ 195,
+ 3027,
+ 2063,
+ 1117,
+ 2473,
+ 4609,
+ 4235,
+ 851,
+ 2719,
+ 3994,
+ 1833,
+ 1485,
+ 558,
+ 4642,
+ 1907,
+ 3934,
+ 4257,
+ 3293,
+ 4631,
+ 3784,
+ 3081,
+ 3931,
+ 644,
+ 2175,
+ 63,
+ 4449,
+ 3152,
+ 1702,
+ 1354,
+ 4622,
+ 3376,
+ 425,
+ 3269,
+ 2016,
+ 2405,
+ 1896,
+ 4243,
+ 2793,
+ 1903,
+ 2256,
+ 1726,
+ 2574,
+ 2877,
+ 3717,
+ 4456,
+ 2476,
+ 784,
+ 4426,
+ 938,
+ 819,
+ 724,
+ 1939,
+ 2795,
+ 179,
+ 3748,
+ 1792,
+ 349,
+ 302,
+ 2128,
+ 3744,
+ 2715,
+ 4491,
+ 1705,
+ 1860,
+ 2740,
+ 2565,
+ 3113,
+ 3661,
+ 2919,
+ 2954,
+ 226,
+ 968,
+ 4019,
+ 1966,
+ 4494,
+ 1402,
+ 4285,
+ 3049,
+ 478,
+ 1920,
+ 4439,
+ 2927,
+ 358,
+ 3316,
+ 1230,
+ 688,
+ 1314,
+ 2386,
+ 862,
+ 1371,
+ 930,
+ 4640,
+ 2088,
+ 629,
+ 3743,
+ 3510,
+ 4056,
+ 1853,
+ 4027,
+ 4432,
+ 3941,
+ 1131,
+ 1692,
+ 2812,
+ 3762,
+ 1565,
+ 4696,
+ 2869,
+ 4041,
+ 3356,
+ 670,
+ 1766,
+ 1147,
+ 2216,
+ 2176,
+ 298,
+ 3087,
+ 282,
+ 1448,
+ 4530,
+ 3144,
+ 4589,
+ 746,
+ 794,
+ 2588,
+ 4469,
+ 4134,
+ 3179,
+ 363,
+ 684,
+ 2000,
+ 1337,
+ 3878,
+ 1194,
+ 3467,
+ 1456,
+ 3680,
+ 2320,
+ 4028,
+ 1646,
+ 3794,
+ 1254,
+ 2815,
+ 651,
+ 4510,
+ 213,
+ 1014,
+ 1652,
+ 3241,
+ 1723,
+ 2874,
+ 3688,
+ 3854,
+ 1205,
+ 828,
+ 4237,
+ 3623,
+ 4520,
+ 1469,
+ 1821,
+ 4278,
+ 2939,
+ 1084,
+ 1842,
+ 2477,
+ 3799,
+ 2933,
+ 2029,
+ 1264,
+ 1166,
+ 1942,
+ 1446,
+ 2697,
+ 2389,
+ 1093,
+ 3755,
+ 1118,
+ 2840,
+ 3190,
+ 484,
+ 3820,
+ 2370,
+ 259,
+ 2463,
+ 2069,
+ 3110,
+ 1964,
+ 2599,
+ 4338,
+ 1904,
+ 3827,
+ 3504,
+ 3751,
+ 655,
+ 113,
+ 80,
+ 3126,
+ 974,
+ 3415,
+ 3528,
+ 2701,
+ 2837,
+ 1959,
+ 2801,
+ 97,
+ 4242,
+ 4537,
+ 4071,
+ 320,
+ 4136,
+ 2356,
+ 3336,
+ 3339,
+ 4658,
+ 3566,
+ 1924,
+ 2947,
+ 3619,
+ 2242,
+ 2475,
+ 2209,
+ 2111,
+ 2208,
+ 1801,
+ 2797,
+ 1751,
+ 4549,
+ 3271,
+ 2800,
+ 3202,
+ 669,
+ 791,
+ 3322,
+ 190,
+ 3745,
+ 1444,
+ 453,
+ 592,
+ 1344,
+ 1957,
+ 4688,
+ 805,
+ 4650,
+ 1808,
+ 3543,
+ 2827,
+ 4472,
+ 482,
+ 4706,
+ 4197,
+ 3806,
+ 3511,
+ 4275,
+ 2351,
+ 3724,
+ 4479,
+ 906,
+ 4233,
+ 3573,
+ 798,
+ 2861,
+ 1632,
+ 2693,
+ 395,
+ 180,
+ 3388,
+ 3811,
+ 840,
+ 305,
+ 2318,
+ 1585,
+ 563,
+ 910,
+ 4506,
+ 578,
+ 788,
+ 1759,
+ 394,
+ 2255,
+ 3942,
+ 3089,
+ 1,
+ 3988,
+ 2010,
+ 4097,
+ 136,
+ 3696,
+ 2092,
+ 1976,
+ 3013,
+ 3147,
+ 2281,
+ 1182,
+ 3288,
+ 841,
+ 3974,
+ 1307,
+ 1786,
+ 631,
+ 1918,
+ 3404,
+ 2745,
+ 1442,
+ 3235,
+ 2678,
+ 3486,
+ 4115,
+ 2068,
+ 424,
+ 3800,
+ 3319,
+ 1956,
+ 2036,
+ 2847,
+ 1681,
+ 3838,
+ 4113,
+ 2689,
+ 3325,
+ 1046,
+ 1804,
+ 4269,
+ 3448,
+ 72,
+ 1516,
+ 1746,
+ 1829,
+ 1135,
+ 1706,
+ 2491,
+ 874,
+ 2239,
+ 2594,
+ 706,
+ 1665,
+ 3993,
+ 2605,
+ 4038,
+ 1281,
+ 1982,
+ 3062,
+ 876,
+ 1331,
+ 703,
+ 2399,
+ 2816,
+ 4310,
+ 1359,
+ 3700,
+ 691,
+ 1546,
+ 2809,
+ 2276,
+ 4425,
+ 3756,
+ 2942,
+ 3633,
+ 253,
+ 2532,
+ 1643,
+ 2379,
+ 3671,
+ 4705,
+ 4442,
+ 3869,
+ 2824,
+ 1938,
+ 3180,
+ 4598,
+ 4307,
+ 1497,
+ 4431,
+ 16,
+ 2373,
+ 672,
+ 418,
+ 1635,
+ 1239,
+ 1573,
+ 3039,
+ 3158,
+ 1366,
+ 1930,
+ 3117,
+ 4185,
+ 4701,
+ 1167,
+ 2771,
+ 1312,
+ 1943,
+ 3465,
+ 76,
+ 4500,
+ 3944,
+ 4619,
+ 3125,
+ 4044,
+ 4675,
+ 3329,
+ 4691,
+ 4637,
+ 1069,
+ 3779,
+ 621,
+ 389,
+ 1929,
+ 1369,
+ 42,
+ 1877,
+ 1970,
+ 3953,
+ 2095,
+ 1802,
+ 1999,
+ 3100,
+ 4054,
+ 3193,
+ 83,
+ 559,
+ 3673,
+ 4416,
+ 4414,
+ 1714,
+ 2245,
+ 2011,
+ 2289,
+ 3456,
+ 4630,
+ 2277,
+ 3466,
+ 4216,
+ 3055,
+ 471,
+ 994,
+ 2118,
+ 1750,
+ 4169,
+ 1655,
+ 3716,
+ 8,
+ 1764,
+ 2892,
+ 1577,
+ 1325,
+ 1973,
+ 3256,
+ 426,
+ 2481,
+ 1630,
+ 2201,
+ 1463,
+ 1606,
+ 1509,
+ 1350,
+ 3786,
+ 3314,
+ 1828,
+ 1578,
+ 641,
+ 4544,
+ 4251,
+ 1986,
+ 3951,
+ 2603,
+ 3634,
+ 2656,
+ 300,
+ 1543,
+ 1180,
+ 4686,
+ 230,
+ 2586,
+ 820,
+ 1372,
+ 1958,
+ 2394,
+ 4571,
+ 4102,
+ 2522,
+ 751,
+ 2974,
+ 2789,
+ 3422,
+ 936,
+ 3006,
+ 3918,
+ 4080,
+ 4248,
+ 1985,
+ 2094,
+ 2580,
+ 2374,
+ 3883,
+ 1908,
+ 1941,
+ 963,
+ 3077,
+ 3426,
+ 4316,
+ 698,
+ 4467,
+ 2238,
+ 2446,
+ 1526,
+ 2534,
+ 1088,
+ 2576,
+ 707,
+ 3782,
+ 2222,
+ 3618,
+ 2421,
+ 809,
+ 1620,
+ 2344,
+ 1002,
+ 104,
+ 4461,
+ 2857,
+ 4200,
+ 1893,
+ 1586,
+ 2097,
+ 4085,
+ 3948,
+ 4611,
+ 879,
+ 3517,
+ 2722,
+ 1298,
+ 2136,
+ 495,
+ 229,
+ 3666,
+ 297,
+ 2746,
+ 2456,
+ 1520,
+ 1445,
+ 941,
+ 2879,
+ 877,
+ 1439,
+ 2778,
+ 2843,
+ 2849,
+ 3736,
+ 3645,
+ 3670,
+ 661,
+ 3833,
+ 3960,
+ 143,
+ 1454,
+ 2292,
+ 575,
+ 4263,
+ 4505,
+ 3385,
+ 3199,
+ 3093,
+ 1290,
+ 3310,
+ 4664,
+ 3228,
+ 2145,
+ 1996,
+ 353,
+ 2687,
+ 404,
+ 2863,
+ 299,
+ 3959,
+ 2870,
+ 1688,
+ 2123,
+ 39,
+ 2169,
+ 1103,
+ 2041,
+ 3183,
+ 1082,
+ 1687,
+ 1857,
+ 2989,
+ 3173,
+ 2972,
+ 1488,
+ 2033,
+ 2072,
+ 1771,
+ 3130,
+ 2488,
+ 2395,
+ 2864,
+ 2297,
+ 4313,
+ 2909,
+ 3411,
+ 3014,
+ 251,
+ 2478,
+ 888,
+ 2984,
+ 102,
+ 1500,
+ 138,
+ 2447,
+ 1263,
+ 1286,
+ 2624,
+ 1790,
+ 3685,
+ 2839,
+ 2414,
+ 4022,
+ 915,
+ 4152,
+ 2770,
+ 2787,
+ 1547,
+ 1011,
+ 4645,
+ 432,
+ 1563,
+ 2458,
+ 928,
+ 71,
+ 573,
+ 4224,
+ 2339,
+ 4497,
+ 528,
+ 3274,
+ 3909,
+ 4410,
+ 1739,
+ 2464,
+ 3272,
+ 4362,
+ 3223,
+ 2151,
+ 789,
+ 4081,
+ 2340,
+ 3879,
+ 2191,
+ 3285,
+ 4229,
+ 1219,
+ 3764,
+ 2854,
+ 3731,
+ 4129,
+ 393,
+ 1429,
+ 2099,
+ 3841,
+ 3019,
+ 386,
+ 4297,
+ 3273,
+ 4119,
+ 4217,
+ 1763,
+ 4680,
+ 524,
+ 1184,
+ 987,
+ 1124,
+ 2224,
+ 1321,
+ 2655,
+ 2926,
+ 2995,
+ 920,
+ 4618,
+ 4366,
+ 4540,
+ 1932,
+ 391,
+ 689,
+ 2681,
+ 1700,
+ 3171,
+ 686,
+ 3191,
+ 699,
+ 3323,
+ 2085,
+ 3290,
+ 2375,
+ 1846,
+ 4514,
+ 1285,
+ 796,
+ 4499,
+ 2686,
+ 4140,
+ 1352,
+ 1060,
+ 4488,
+ 3028,
+ 4657,
+ 2792,
+ 2545,
+ 4379,
+ 2621,
+ 934,
+ 3098,
+ 3725,
+ 3689,
+ 1725,
+ 4030,
+ 3221,
+ 4173,
+ 374,
+ 2951,
+ 2451,
+ 1748,
+ 2790,
+ 3240,
+ 4009,
+ 3753,
+ 1007,
+ 721,
+ 1740,
+ 4623,
+ 958,
+ 2416,
+ 1110,
+ 3253,
+ 1864,
+ 1789,
+ 3981,
+ 2873,
+ 1087,
+ 4582,
+ 3102,
+ 1271,
+ 1072,
+ 3080,
+ 1215,
+ 2564,
+ 579,
+ 3211,
+ 1605,
+ 74,
+ 2561,
+ 980,
+ 4205,
+ 650,
+ 3279,
+ 3582,
+ 4679,
+ 3713,
+ 2079,
+ 2330,
+ 3651,
+ 2159,
+ 4573,
+ 3958,
+ 2807,
+ 3663,
+ 2774,
+ 1291,
+ 2024,
+ 1496,
+ 1409,
+ 523,
+ 3107,
+ 158,
+ 9,
+ 513,
+ 3461,
+ 3216,
+ 1735,
+ 4268,
+ 60,
+ 2165,
+ 3649,
+ 3665,
+ 1081,
+ 3491,
+ 4180,
+ 1809,
+ 1390,
+ 2741,
+ 3134,
+ 2294,
+ 1278,
+ 4279,
+ 887,
+ 4511,
+ 84,
+ 218,
+ 2525,
+ 919,
+ 1003,
+ 3857,
+ 3073,
+ 414,
+ 2734,
+ 979,
+ 562,
+ 1962,
+ 4330,
+ 3868,
+ 4460,
+ 992,
+ 2398,
+ 4223,
+ 624,
+ 1034,
+ 3420,
+ 1887,
+ 165,
+ 93,
+ 103,
+ 2156,
+ 818,
+ 3056,
+ 893,
+ 1867,
+ 1875,
+ 2783,
+ 3210,
+ 3032,
+ 3133,
+ 2766,
+ 1922,
+ 1720,
+ 1071,
+ 3824,
+ 2965,
+ 1225,
+ 2127,
+ 2415,
+ 65,
+ 4020,
+ 2492,
+ 4614,
+ 4281,
+ 4548,
+ 3335,
+ 3899,
+ 3604,
+ 924,
+ 304,
+ 1732,
+ 3277,
+ 3150,
+ 2101,
+ 2673,
+ 3225,
+ 437,
+ 3333,
+ 2531,
+ 1317,
+ 4012,
+ 2796,
+ 4077,
+ 1933,
+ 4144,
+ 264,
+ 3704,
+ 3116,
+ 3384,
+ 2910,
+ 743,
+ 128,
+ 3071,
+ 480,
+ 3205,
+ 2997,
+ 4333,
+ 2455,
+ 338,
+ 3945,
+ 666,
+ 2539,
+ 2996,
+ 3720,
+ 111,
+ 2763,
+ 4597,
+ 3001,
+ 3234,
+ 2994,
+ 3954,
+ 825,
+ 4396,
+ 3678,
+ 3108,
+ 3230,
+ 2179,
+ 3962,
+ 2114,
+ 1506,
+ 4369,
+ 519,
+ 4608,
+ 834,
+ 4417,
+ 3710,
+ 3255,
+ 3427,
+ 3229,
+ 4665,
+ 3547,
+ 19,
+ 3997,
+ 1529,
+ 1691,
+ 1332,
+ 3470,
+ 2160,
+ 2148,
+ 873,
+ 154,
+ 1015,
+ 2526,
+ 3474,
+ 4018,
+ 122,
+ 2411,
+ 2272,
+ 2904,
+ 2779,
+ 1475,
+ 3311,
+ 2844,
+ 1195,
+ 2806,
+ 1164,
+ 2999,
+ 232,
+ 3050,
+ 3886,
+ 748,
+ 1768,
+ 2623,
+ 3501,
+ 1218,
+ 4298,
+ 4201,
+ 4003,
+ 1246,
+ 3772,
+ 4176,
+ 2346,
+ 1201,
+ 3189,
+ 1383,
+ 4551,
+ 243,
+ 2210,
+ 3186,
+ 1602,
+ 14,
+ 3992,
+ 871,
+ 4590,
+ 1514,
+ 2290,
+ 758,
+ 3481,
+ 2762,
+ 4339,
+ 952,
+ 4628,
+ 3118,
+ 167,
+ 1507,
+ 1216,
+ 4448,
+ 2511,
+ 2834,
+ 2846,
+ 3209,
+ 4105,
+ 2215,
+ 1512,
+ 3776,
+ 4095,
+ 3406,
+ 3451,
+ 636,
+ 1703,
+ 172,
+ 2611,
+ 3166,
+ 3246,
+ 1262,
+ 3662,
+ 1424,
+ 2543,
+ 1569,
+ 2585,
+ 3220,
+ 1713,
+ 1978,
+ 1270,
+ 1910,
+ 1381,
+ 4063,
+ 728,
+ 4074,
+ 2352,
+ 2316,
+ 1731,
+ 3735,
+ 3478,
+ 3520,
+ 286,
+ 4668,
+ 4354,
+ 744,
+ 3576,
+ 3906,
+ 676,
+ 1532,
+ 2450,
+ 416,
+ 949,
+ 892,
+ 656,
+ 1716,
+ 927,
+ 1876,
+ 3503,
+ 3483,
+ 1889,
+ 3853,
+ 1214,
+ 3654,
+ 2494,
+ 3161,
+ 3846,
+ 2396,
+ 473,
+ 3487,
+ 1052,
+ 3387,
+ 2189,
+ 1280,
+ 3402,
+ 2998,
+ 2496,
+ 58,
+ 1373,
+ 106,
+ 1510,
+ 3551,
+ 3803,
+ 3669,
+ 1258,
+ 648,
+ 4495,
+ 1848,
+ 1580,
+ 2501,
+ 3114,
+ 1259,
+ 1812,
+ 3600,
+ 3364,
+ 3480,
+ 2115,
+ 3489,
+ 303,
+ 739,
+ 361,
+ 802,
+ 2368,
+ 192,
+ 1028,
+ 1601,
+ 2265,
+ 4415,
+ 326,
+ 1972,
+ 4512,
+ 1788,
+ 3473,
+ 3300,
+ 1979,
+ 1056,
+ 1648,
+ 4711,
+ 2973,
+ 4585,
+ 4478,
+ 3590,
+ 479,
+ 1169,
+ 182,
+ 534,
+ 3983,
+ 4182,
+ 140,
+ 1594,
+ 1186,
+ 4629,
+ 2731,
+ 1106,
+ 3967,
+ 3157,
+ 4434,
+ 315,
+ 2589,
+ 2490,
+ 2946,
+ 1311,
+ 860,
+ 1795,
+ 489,
+ 908,
+ 901,
+ 4684,
+ 1631,
+ 4647,
+ 780,
+ 137,
+ 2039,
+ 2513,
+ 3317,
+ 1744,
+ 1473,
+ 1063,
+ 1461,
+ 550,
+ 4300,
+ 2584,
+ 3818,
+ 3296,
+ 549,
+ 3007,
+ 2845,
+ 2229,
+ 2546,
+ 412,
+ 4504,
+ 2744,
+ 2728,
+ 2964,
+ 2842,
+ 2704,
+ 4464,
+ 217,
+ 4490,
+ 3418,
+ 2296,
+ 1145,
+ 2194,
+ 2467,
+ 922,
+ 2130,
+ 1389,
+ 1554,
+ 4303,
+ 2684,
+ 3372,
+ 4084,
+ 2205,
+ 1535,
+ 3536,
+ 4559,
+ 859,
+ 3943,
+ 440,
+ 2712,
+ 1370,
+ 1363,
+ 3990,
+ 4433,
+ 2217,
+ 2661,
+ 2608,
+ 723,
+ 2190,
+ 1961,
+ 1798,
+ 2757,
+ 4681,
+ 981,
+ 2383,
+ 4040,
+ 3337,
+ 187,
+ 2601,
+ 3366,
+ 3817,
+ 2044,
+ 15,
+ 2034,
+ 4092,
+ 3578,
+ 397,
+ 4389,
+ 1163,
+ 2,
+ 3115,
+ 1905,
+ 3312,
+ 1368,
+ 772,
+ 2880,
+ 376,
+ 4091,
+ 986,
+ 4262,
+ 3509,
+ 1712,
+ 3453,
+ 4013,
+ 1880,
+ 2575,
+ 3788,
+ 1293,
+ 867,
+ 2889,
+ 3188,
+ 764,
+ 3734,
+ 4513,
+ 2349,
+ 1336,
+ 2577,
+ 2515,
+ 4334,
+ 3381,
+ 2716,
+ 4451,
+ 3529,
+ 756,
+ 1600,
+ 1070,
+ 1738,
+ 790,
+ 3652,
+ 1414,
+ 3726,
+ 1523,
+ 4127,
+ 1247,
+ 921,
+ 422,
+ 3605,
+ 1356,
+ 2445,
+ 199,
+ 1521,
+ 41,
+ 1417,
+ 1925,
+ 4343,
+ 1361,
+ 1155,
+ 2924,
+ 1803,
+ 2443,
+ 1422,
+ 3923,
+ 1148,
+ 811,
+ 4138,
+ 2991,
+ 3392,
+ 4534,
+ 3010,
+ 378,
+ 78,
+ 4032,
+ 110,
+ 778,
+ 1579,
+ 3302,
+ 3610,
+ 3022,
+ 3286,
+ 3580,
+ 3613,
+ 984,
+ 354,
+ 2969,
+ 4427,
+ 4282,
+ 245,
+ 275,
+ 3506,
+ 183,
+ 444,
+ 2231,
+ 1413,
+ 2059,
+ 880,
+ 263,
+ 1207,
+ 1272,
+ 2538,
+ 3887,
+ 1662,
+ 4620,
+ 2784,
+ 3320,
+ 2956,
+ 3999,
+ 2593,
+ 4293,
+ 3609,
+ 2497,
+ 900,
+ 3643,
+ 225,
+ 3142,
+ 637,
+ 4600,
+ 3664,
+ 3003,
+ 4037,
+ 3963,
+ 933,
+ 6,
+ 3719,
+ 2516,
+ 4667,
+ 4067,
+ 800,
+ 169,
+ 613,
+ 2620,
+ 3018,
+ 457,
+ 3140,
+ 2976,
+ 1633,
+ 1115,
+ 1160,
+ 3783,
+ 4162,
+ 3874,
+ 1610,
+ 1855,
+ 543,
+ 472,
+ 4639,
+ 1722,
+ 2212,
+ 3493,
+ 1729,
+ 3280,
+ 1541,
+ 1616,
+ 2113,
+ 4344,
+ 2634,
+ 2030,
+ 1858,
+ 3569,
+ 2248,
+ 3819,
+ 3195,
+ 589,
+ 2759,
+ 725,
+ 2214,
+ 2162,
+ 902,
+ 904,
+ 1742,
+ 166,
+ 4314,
+ 3810,
+ 1236,
+ 2312,
+ 3858,
+ 3362,
+ 2558,
+ 940,
+ 3298,
+ 4204,
+ 2021,
+ 540,
+ 2082,
+ 1176,
+ 1830,
+ 2658,
+ 455,
+ 3419,
+ 3127,
+ 4159,
+ 508,
+ 3500,
+ 4326,
+ 2609,
+ 1612,
+ 2378,
+ 935,
+ 822,
+ 434,
+ 890,
+ 4421,
+ 148,
+ 2064,
+ 3922,
+ 4114,
+ 546,
+ 1426,
+ 2707,
+ 554,
+ 1709,
+ 2723,
+ 1753,
+ 1940,
+ 806,
+ 2597,
+ 1345,
+ 2357,
+ 4052,
+ 4327,
+ 356,
+ 3058,
+ 625,
+ 77,
+ 4452,
+ 151,
+ 2918,
+ 2419,
+ 2671,
+ 2484,
+ 1441,
+ 1235,
+ 367,
+ 321,
+ 1909,
+ 1668,
+ 3521,
+ 4526,
+ 3472,
+ 3512,
+ 3864,
+ 4536,
+ 693,
+ 57,
+ 4699,
+ 598,
+ 3956,
+ 3787,
+ 1326,
+ 3773,
+ 1355,
+ 1136,
+ 342,
+ 4096,
+ 3238,
+ 3917,
+ 1562,
+ 2721,
+ 4708,
+ 202,
+ 3160,
+ 1883,
+ 2350,
+ 3587,
+ 2440,
+ 3386,
+ 48,
+ 2121,
+ 4221,
+ 2548,
+ 4351,
+ 4659,
+ 2569,
+ 3450,
+ 3946,
+ 4234,
+ 4031,
+ 2073,
+ 1300,
+ 341,
+ 452,
+ 1315,
+ 2426,
+ 176,
+ 1348,
+ 866,
+ 505,
+ 1025,
+ 4378,
+ 846,
+ 1038,
+ 2441,
+ 1710,
+ 2542,
+ 2358,
+ 4346,
+ 3920,
+ 2112,
+ 792,
+ 2066,
+ 3822,
+ 526,
+ 2149,
+ 3005,
+ 4455,
+ 3033,
+ 872,
+ 3,
+ 2436,
+ 3905,
+ 3287,
+ 231,
+ 3832,
+ 1504,
+ 3660,
+ 3194,
+ 3849,
+ 785,
+ 4621,
+ 1525,
+ 3703,
+ 4329,
+ 2875,
+ 4073,
+ 4186,
+ 2420,
+ 4632,
+ 1882,
+ 2649,
+ 4374,
+ 3985,
+ 4360,
+ 2810,
+ 2323,
+ 1269,
+ 383,
+ 4532,
+ 1287,
+ 1927,
+ 3859,
+ 2153,
+ 704,
+ 634,
+ 4615,
+ 1157,
+ 577,
+ 4252,
+ 4157,
+ 4669,
+ 3398,
+ 964,
+ 2173,
+ 1589,
+ 1055,
+ 4211,
+ 1639,
+ 1895,
+ 3401,
+ 1743,
+ 3423,
+ 2591,
+ 730,
+ 1773,
+ 1059,
+ 2720,
+ 869,
+ 4057,
+ 2305,
+ 726,
+ 3694,
+ 4132,
+ 486,
+ 595,
+ 1666,
+ 308,
+ 3586,
+ 2695,
+ 500,
+ 1934,
+ 387,
+ 3124,
+ 3182,
+ 4613,
+ 2269,
+ 2371,
+ 3072,
+ 545,
+ 830,
+ 2137,
+ 1253,
+ 1004,
+ 4588,
+ 3554,
+ 2003,
+ 745,
+ 4547,
+ 3120,
+ 4443,
+ 3090,
+ 923,
+ 1376,
+ 1420,
+ 3975,
+ 1576,
+ 1030,
+ 4578,
+ 272,
+ 3031,
+ 4682,
+ 4286,
+ 4283,
+ 3477,
+ 1482,
+ 3000,
+ 1464,
+ 2941,
+ 3518,
+ 351,
+ 319,
+ 295,
+ 4306,
+ 1856,
+ 2894,
+ 1349,
+ 125,
+ 1162,
+ 529,
+ 4440,
+ 3260,
+ 757,
+ 2659,
+ 4403,
+ 2718,
+ 2767,
+ 3749,
+ 372,
+ 101,
+ 2207,
+ 3168,
+ 2557,
+ 4355,
+ 814,
+ 2814,
+ 2327,
+ 899,
+ 343,
+ 1251,
+ 4321,
+ 2958,
+ 889,
+ 1840,
+ 3845,
+ 3914,
+ 2982,
+ 4123,
+ 2803,
+ 4489,
+ 3532,
+ 2662,
+ 2226,
+ 1126,
+ 3067,
+ 2090,
+ 1421,
+ 450,
+ 4413,
+ 1096,
+ 61,
+ 2462,
+ 539,
+ 3758,
+ 3677,
+ 4492,
+ 2832,
+ 3245,
+ 2393,
+ 850,
+ 2449,
+ 428,
+ 1664,
+ 476,
+ 4445,
+ 4189,
+ 1343,
+ 1990,
+ 3565,
+ 835,
+ 3933,
+ 690,
+ 956,
+ 3508,
+ 3295,
+ 186,
+ 2878,
+ 4006,
+ 3555,
+ 3599,
+ 2091,
+ 2108,
+ 4438,
+ 2666,
+ 2282,
+ 4155,
+ 749,
+ 2683,
+ 3309,
+ 2388,
+ 307,
+ 4021,
+ 884,
+ 240,
+ 2120,
+ 3218,
+ 3023,
+ 4296,
+ 3907,
+ 3796,
+ 1553,
+ 2743,
+ 4648,
+ 1805,
+ 2817,
+ 3697,
+ 970,
+ 4273,
+ 2626,
+ 3769,
+ 1198,
+ 674,
+ 3602,
+ 2461,
+ 3686,
+ 2619,
+ 1431,
+ 2510,
+ 628,
+ 3642,
+ 913,
+ 3754,
+ 3196,
+ 227,
+ 1675,
+ 2667,
+ 3454,
+ 3263,
+ 2381,
+ 2617,
+ 1741,
+ 1581,
+ 1178,
+ 17,
+ 2906,
+ 700,
+ 1721,
+ 1890,
+ 3538,
+ 1797,
+ 105,
+ 3341,
+ 1261,
+ 3791,
+ 4220,
+ 13,
+ 685,
+ 2438,
+ 3856,
+ 3996,
+ 2058,
+ 2835,
+ 3266,
+ 2423,
+ 3928,
+ 4335,
+ 1844,
+ 1551,
+ 2045,
+ 4167,
+ 4118,
+ 1873,
+ 3358,
+ 1825,
+ 2508,
+ 1954,
+ 1760,
+ 3620,
+ 3377,
+ 1822,
+ 2172,
+ 4148,
+ 3297,
+ 3141,
+ 194,
+ 680,
+ 1303,
+ 2400,
+ 1347,
+ 3542,
+ 2054,
+ 878,
+ 4601,
+ 3658,
+ 817,
+ 626,
+ 2914,
+ 3693,
+ 2047,
+ 3249,
+ 86,
+ 4082,
+ 1137,
+ 1041,
+ 3175,
+ 1342,
+ 2897,
+ 99,
+ 3888,
+ 1029,
+ 4409,
+ 212,
+ 1608,
+ 3409,
+ 3065,
+ 4407,
+ 1595,
+ 566,
+ 365,
+ 1619,
+ 4400,
+ 3656,
+ 2674,
+ 2813,
+ 512,
+ 587,
+ 2694,
+ 2777,
+ 345,
+ 3099,
+ 2184,
+ 3884,
+ 3594,
+ 675,
+ 1659,
+ 1993,
+ 1120,
+ 713,
+ 4695,
+ 2987,
+ 2397,
+ 2380,
+ 3986,
+ 1814,
+ 2407,
+ 1211,
+ 348,
+ 2696,
+ 3955,
+ 3393,
+ 732,
+ 3405,
+ 1900,
+ 1676,
+ 334,
+ 2454,
+ 400,
+ 705,
+ 1375,
+ 1240,
+ 2890,
+ 2709,
+ 1952,
+ 2119,
+ 3303,
+ 895,
+ 2086,
+ 1641,
+ 3442,
+ 2752,
+ 2753,
+ 774,
+ 608,
+ 2359,
+ 3624,
+ 2703,
+ 3978,
+ 865,
+ 4524,
+ 4325,
+ 883,
+ 3201,
+ 2078,
+ 551,
+ 2782,
+ 612,
+ 1204,
+ 2195,
+ 3523,
+ 1384,
+ 1820,
+ 3380,
+ 3363,
+ 25,
+ 3267,
+ 4000,
+ 2080,
+ 2135,
+ 81,
+ 2268,
+ 1680,
+ 3122,
+ 1210,
+ 3839,
+ 4337,
+ 4246,
+ 3395,
+ 1636,
+ 2211,
+ 1466,
+ 442,
+ 4258,
+ 844,
+ 3564,
+ 3885,
+ 4465,
+ 3835,
+ 381,
+ 2859,
+ 2427,
+ 3991,
+ 285,
+ 3679,
+ 2729,
+ 4666,
+ 4043,
+ 2452,
+ 555,
+ 1132,
+ 4562,
+ 2614,
+ 2950,
+ 242,
+ 130,
+ 959,
+ 2639,
+ 1127,
+ 781,
+ 4700,
+ 3447,
+ 1708,
+ 2404,
+ 3650,
+ 3667,
+ 4678,
+ 1826,
+ 2944,
+ 3706,
+ 4340,
+ 2197,
+ 3276,
+ 2596,
+ 3682,
+ 278,
+ 885,
+ 3265,
+ 2302,
+ 4388,
+ 2578,
+ 2675,
+ 2345,
+ 1931,
+ 1963,
+ 1658,
+ 3681,
+ 147,
+ 4109,
+ 1737,
+ 1152,
+ 1221,
+ 768,
+ 829,
+ 3957,
+ 4312,
+ 603,
+ 2457,
+ 4117,
+ 1452,
+ 2622,
+ 2171,
+ 1175,
+ 2710,
+ 3527,
+ 2122,
+ 3715,
+ 3318,
+ 1917,
+ 2141,
+ 1823,
+ 2052,
+ 4428,
+ 783,
+ 2096,
+ 2650,
+ 340,
+ 174,
+ 1158,
+ 2218,
+ 4260,
+ 4090,
+ 3766,
+ 284,
+ 717,
+ 531,
+ 3092,
+ 863,
+ 868,
+ 978,
+ 3526,
+ 157,
+ 803,
+ 2798,
+ 1921,
+ 2116,
+ 4236,
+ 3834,
+ 2309,
+ 1128,
+ 4163,
+ 369,
+ 1881,
+ 960,
+ 2872,
+ 2657,
+ 66,
+ 1334,
+ 813,
+ 565,
+ 1224,
+ 4516,
+ 947,
+ 267,
+ 447,
+ 3940,
+ 2050,
+ 2465,
+ 1852,
+ 4685,
+ 1418,
+ 4060,
+ 129,
+ 82,
+ 903,
+ 1244,
+ 2644,
+ 4357,
+ 1936,
+ 2628,
+ 64,
+ 2431,
+ 1172,
+ 4617,
+ 4214,
+ 3069,
+ 2968,
+ 2146,
+ 1831,
+ 3035,
+ 4181,
+ 638,
+ 3440,
+ 1249,
+ 3162,
+ 1212,
+ 1451,
+ 948,
+ 3979,
+ 1209,
+ 1872,
+ 590,
+ 4625,
+ 2641,
+ 3350,
+ 2826,
+ 1220,
+ 1050,
+ 3875,
+ 256,
+ 2267,
+ 2676,
+ 3525,
+ 3226,
+ 3138,
+ 3571,
+ 4674,
+ 4244,
+ 1040,
+ 4368,
+ 311,
+ 2960,
+ 3950,
+ 1975,
+ 152,
+ 833,
+ 552,
+ 2848,
+ 2631,
+ 4517,
+ 2907,
+ 3457,
+ 3640,
+ 3562,
+ 2917,
+ 931,
+ 1396,
+ 3770,
+ 3370,
+ 1273,
+ 368,
+ 1401,
+ 4627,
+ 1984,
+ 2203,
+ 1265,
+ 1779,
+ 773,
+ 564,
+ 4457,
+ 4253,
+ 955,
+ 3270,
+ 1673,
+ 1134,
+ 701,
+ 1545,
+ 3070,
+ 2761,
+ 1694,
+ 3659,
+ 3136,
+ 465,
+ 7,
+ 290,
+ 782,
+ 731,
+ 982,
+ 1140,
+ 4294,
+ 134,
+ 446,
+ 3598,
+ 622,
+ 289,
+ 339,
+ 2922,
+ 1967,
+ 3847,
+ 3106,
+ 1552,
+ 1433,
+ 3294,
+ 679,
+ 594,
+ 4299,
+ 643,
+ 673,
+ 2529,
+ 4068,
+ 322,
+ 3574,
+ 327,
+ 123,
+ 568,
+ 4121,
+ 3416,
+ 793,
+ 1010,
+ 4116,
+ 3641,
+ 483,
+ 3494,
+ 2071,
+ 2360,
+ 1609,
+ 3935,
+ 107,
+ 1111,
+ 4110,
+ 4295,
+ 4219,
+ 1470,
+ 135,
+ 4142,
+ 1783,
+ 144,
+ 3176,
+ 708,
+ 2313,
+ 696,
+ 1434,
+ 38,
+ 3378,
+ 2170,
+ 1415,
+ 4072,
+ 1950,
+ 2026,
+ 1682,
+ 600,
+ 4535,
+ 2186,
+ 2022,
+ 37,
+ 1564,
+ 2315,
+ 3304,
+ 1806,
+ 2042,
+ 1862,
+ 4653,
+ 1916,
+ 3890,
+ 4371,
+ 1139,
+ 2437,
+ 3672,
+ 3020,
+ 1091,
+ 1054,
+ 2888,
+ 2430,
+ 1100,
+ 1915,
+ 23,
+ 2012,
+ 4422,
+ 1043,
+ 4227,
+ 3131,
+ 4676,
+ 126,
+ 1447,
+ 2648,
+ 2652,
+ 1678,
+ 4290,
+ 4587,
+ 1200,
+ 1020,
+ 3977,
+ 4172,
+ 2278,
+ 1838,
+ 2070,
+ 2627,
+ 1518,
+ 2004,
+ 1292,
+ 1149,
+ 3575,
+ 1367,
+ 3881,
+ 3552,
+ 313,
+ 727,
+ 2544,
+ 2480,
+ 2867,
+ 1479,
+ 439,
+ 2637,
+ 821,
+ 627,
+ 98,
+ 1926,
+ 1123,
+ 1129,
+ 3777,
+ 1997,
+ 3455,
+ 1459,
+ 4703,
+ 4048,
+ 4207,
+ 2841,
+ 4480,
+ 94,
+ 2286,
+ 433,
+ 3626,
+ 456,
+ 1185,
+ 2902,
+ 4014,
+ 3469,
+ 4150,
+ 3814,
+ 2325,
+ 3638,
+ 914,
+ 3284,
+ 1513,
+ 945,
+ 2015,
+ 4654,
+ 2571,
+ 2498,
+ 2573,
+ 1231,
+ 2732,
+ 1642,
+ 207,
+ 3752,
+ 1257,
+ 1453,
+ 3608,
+ 3964,
+ 946,
+ 916,
+ 2304,
+ 4564,
+ 4284,
+ 4363,
+ 1484,
+ 3683,
+ 2896,
+ 4274,
+ 4663,
+ 907,
+ 4606,
+ 639,
+ 553,
+ 998,
+ 606,
+ 220,
+ 3326,
+ 740,
+ 2469,
+ 4046,
+ 3894,
+ 4545,
+ 3021,
+ 3495,
+ 816,
+ 4,
+ 4525,
+ 2219,
+ 1133,
+ 1837,
+ 1460,
+ 3059,
+ 2262,
+ 1335,
+ 1092,
+ 2341,
+ 3231,
+ 3399,
+ 1228,
+ 3718,
+ 4550,
+ 3139,
+ 3498,
+ 3904,
+ 1468,
+ 3519,
+ 3502,
+ 3741,
+ 2725,
+ 3177,
+ 532,
+ 2274,
+ 2472,
+ 3254,
+ 3774,
+ 4385,
+ 168,
+ 2384,
+ 2568,
+ 4392,
+ 3400,
+ 1478,
+ 2495,
+ 4697,
+ 3630,
+ 4287,
+ 402,
+ 2764,
+ 2428,
+ 4377,
+ 3002,
+ 4649,
+ 4280,
+ 4198,
+ 2261,
+ 1432,
+ 238,
+ 1645,
+ 3805,
+ 3631
+ ],
+ "val_indices": [
+ 1051,
+ 849,
+ 4437,
+ 847,
+ 738,
+ 1171,
+ 4322,
+ 1476,
+ 1701,
+ 3407,
+ 2887,
+ 3438,
+ 3119,
+ 1365,
+ 4383,
+ 3170,
+ 2257,
+ 3024,
+ 2607,
+ 3146,
+ 3757,
+ 1430,
+ 265,
+ 3247,
+ 468,
+ 1064,
+ 1467,
+ 415,
+ 2157,
+ 2993,
+ 239,
+ 2970,
+ 1412,
+ 3197,
+ 4087,
+ 2084,
+ 1781,
+ 2301,
+ 2541,
+ 4604,
+ 2509,
+ 1827,
+ 3354,
+ 1408,
+ 1104,
+ 881,
+ 1865,
+ 2333,
+ 1623,
+ 4592,
+ 1013,
+ 2642,
+ 1250,
+ 3522,
+ 3424,
+ 177,
+ 1879,
+ 3198,
+ 1075,
+ 1275,
+ 2338,
+ 1946,
+ 1622,
+ 2981,
+ 965,
+ 50,
+ 957,
+ 918,
+ 4039,
+ 4034,
+ 2977,
+ 1191,
+ 1818,
+ 1458,
+ 3707,
+ 2895,
+ 2819,
+ 1296,
+ 729,
+ 527,
+ 3617,
+ 4474,
+ 163,
+ 3507,
+ 3291,
+ 52,
+ 2714,
+ 3111,
+ 2412,
+ 2434,
+ 1505,
+ 90,
+ 838,
+ 861,
+ 645,
+ 3128,
+ 3261,
+ 2424,
+ 1752,
+ 1749,
+ 2618,
+ 4712,
+ 3490,
+ 2326,
+ 4017,
+ 4168,
+ 823,
+ 926,
+ 4226,
+ 448,
+ 2347,
+ 596,
+ 388,
+ 1524,
+ 2751,
+ 2528,
+ 642,
+ 4026,
+ 1309,
+ 1624,
+ 2983,
+ 3949,
+ 2213,
+ 4554,
+ 4156,
+ 4086,
+ 487,
+ 12,
+ 3561,
+ 355,
+ 3976,
+ 3349,
+ 4387,
+ 1308,
+ 4049,
+ 2193,
+ 1073,
+ 2505,
+ 4633,
+ 3345,
+ 3952,
+ 4470,
+ 1511,
+ 2138,
+ 4453,
+ 206,
+ 1318,
+ 3546,
+ 3250,
+ 359,
+ 2077,
+ 2530,
+ 4644,
+ 3908,
+ 1279,
+ 4288,
+ 4382,
+ 268,
+ 810,
+ 1353,
+ 716,
+ 4521,
+ 1960,
+ 808,
+ 43,
+ 1841,
+ 1480,
+ 1374,
+ 2967,
+ 3178,
+ 1031,
+ 1385,
+ 3352,
+ 1533,
+ 1813,
+ 4602,
+ 1892,
+ 2692,
+ 2435,
+ 3622,
+ 516,
+ 616,
+ 2651,
+ 2051,
+ 1114,
+ 1669,
+ 3798,
+ 4528,
+ 4471,
+ 4149,
+ 2663,
+ 3767,
+ 1021,
+ 117,
+ 4656,
+ 2882,
+ 514,
+ 662,
+ 2056,
+ 3545,
+ 1428,
+ 2551,
+ 266,
+ 1302,
+ 175,
+ 2566,
+ 4599,
+ 3078,
+ 3695,
+ 891,
+ 1316,
+ 1816,
+ 1256,
+ 469,
+ 2482,
+ 1425,
+ 4146,
+ 2776,
+ 2852,
+ 839,
+ 121,
+ 4240,
+ 406,
+ 1835,
+ 2303,
+ 3008,
+ 3676,
+ 2332,
+ 4561,
+ 799,
+ 1388,
+ 2038,
+ 588,
+ 2905,
+ 1324,
+ 2107,
+ 1911,
+ 2823,
+ 3367,
+ 1765,
+ 3687,
+ 3169,
+ 2252,
+ 109,
+ 246,
+ 1377,
+ 683,
+ 2273,
+ 1058,
+ 4036,
+ 1018,
+ 4047,
+ 2604,
+ 1888,
+ 2980,
+ 4694,
+ 490,
+ 3655,
+ 3987,
+ 4079,
+ 3075,
+ 3360,
+ 759,
+ 2562,
+ 1607,
+ 3902,
+ 1598,
+ 26,
+ 1398,
+ 995,
+ 2682,
+ 4652,
+ 3563,
+ 2708,
+ 2032,
+ 1006,
+ 2512,
+ 462,
+ 1487,
+ 1121,
+ 3971,
+ 3371,
+ 350,
+ 1522,
+ 3074,
+ 3657,
+ 3727,
+ 2885,
+ 4361,
+ 249,
+ 770,
+ 2263,
+ 1053,
+ 4264,
+ 4107,
+ 2786,
+ 4406,
+ 2567,
+ 3919,
+ 2328,
+ 2632,
+ 385,
+ 3932,
+ 4305,
+ 3584,
+ 557,
+ 3723,
+ 1177,
+ 5,
+ 3143,
+ 1234,
+ 668,
+ 3898,
+ 318,
+ 2961,
+ 399,
+ 1346,
+ 2280,
+ 2500,
+ 4317,
+ 3268,
+ 682,
+ 2319,
+ 3340,
+ 1686,
+ 3698,
+ 2372,
+ 2940,
+ 752,
+ 1306,
+ 2750,
+ 2829,
+ 1884,
+ 937,
+ 294,
+ 3939,
+ 4603,
+ 1777,
+ 4376,
+ 4404,
+ 1437,
+ 2937,
+ 1657,
+ 2417,
+ 3801,
+ 3187,
+ 4581,
+ 3497,
+ 3926,
+ 4576,
+ 371,
+ 1517,
+ 2192,
+ 2876,
+ 3029,
+ 2702,
+ 188,
+ 1130,
+ 1649,
+ 4218,
+ 3275,
+ 274,
+ 55,
+ 3313,
+ 2439,
+ 3132,
+ 1024,
+ 2204,
+ 2978,
+ 1834,
+ 1819,
+ 4375,
+ 4477,
+ 632,
+ 857,
+ 1621,
+ 1679,
+ 2266,
+ 2413,
+ 2821,
+ 3568,
+ 2230,
+ 779,
+ 2699,
+ 4210,
+ 1697,
+ 347,
+ 3242,
+ 2600,
+ 1611,
+ 241,
+ 2314,
+ 1080,
+ 2234,
+ 2425,
+ 4498,
+ 1969,
+ 3446,
+ 3204,
+ 3066,
+ 2006,
+ 4458,
+ 3768,
+ 2527,
+ 3567,
+ 2742,
+ 3040,
+ 3793,
+ 3042,
+ 3289,
+ 1102,
+ 1637,
+ 3930,
+ 4575,
+ 1592,
+ 430,
+ 3530,
+ 4222,
+ 1472,
+ 1410,
+ 3025,
+ 753,
+ 3595,
+ 4436,
+ 1615,
+ 4011,
+ 1301,
+ 3583,
+ 2858,
+ 3867,
+ 2613,
+ 3732,
+ 4475,
+ 2713,
+ 4670,
+ 2772,
+ 421,
+ 3828,
+ 3429,
+ 2945,
+ 3837,
+ 2336,
+ 3980,
+ 1379,
+ 2664,
+ 1206,
+ 4655,
+ 3086,
+ 4565,
+ 1528,
+ 1799,
+ 438,
+ 3305,
+ 2625,
+ 2727,
+ 3129,
+ 3537,
+ 660,
+ 4557,
+ 3921,
+ 826,
+ 509,
+ 3807,
+ 3109,
+ 4709,
+ 2198,
+ 1574,
+ 454,
+ 4501,
+ 3606,
+ 1156
+ ]
+}
diff --git a/property_prediction/configs/spherenet/README.md b/property_prediction/configs/spherenet/README.md
new file mode 100644
index 00000000..ea5c6584
--- /dev/null
+++ b/property_prediction/configs/spherenet/README.md
@@ -0,0 +1,261 @@
+# SphereNet
+
+[Spherical Message Passing for 3D Molecular Graphs](https://arxiv.org/abs/2102.05013) (ICLR 2021)
+
+## Abstract
+
+We propose the spherical message passing (SMP) scheme for 3D molecular graphs,
+which leverages **distance, angle, and torsion** information simultaneously to
+uniquely identify the relative positions of atoms in 3D space. Previous
+methods such as SchNet (distance-only) and DimeNet++ (distance + angle) suffer
+from equivariance ambiguity because multiple spatial configurations can map
+to the same pairwise distances or angles. By incorporating torsion angles
+(dihedral angles), SphereNet resolves this ambiguity and achieves
+state-of-the-art results on the QM9 benchmark.
+
+
+