├── LICENSE ├── README.md ├── datasets └── __init__.py ├── docs ├── changelog.md ├── env_setup.md └── equiformer_v2_poster.pdf ├── engine.py ├── env └── env_equiformer_v2.yml ├── fig ├── equiformer_v2.png ├── equiformer_v2_adsorbml_results.png ├── equiformer_v2_oc20_results.png ├── equiformer_v2_oc22_results.png └── equiformer_v2_speed_accuracy_tradeoffs.png ├── logger.py ├── main_oc20.py ├── nets ├── __init__.py └── equiformer_v2 │ ├── Jd.pt │ ├── activation.py │ ├── drop.py │ ├── edge_rot_mat.py │ ├── equiformer_v2_oc20.py │ ├── gaussian_rbf.py │ ├── input_block.py │ ├── layer_norm.py │ ├── module_list.py │ ├── radial_function.py │ ├── so2_ops.py │ ├── so3.py │ ├── transformer_block.py │ └── wigner.py ├── oc20 ├── configs │ └── s2ef │ │ ├── 2M │ │ ├── base.yml │ │ └── equiformer_v2 │ │ │ ├── equiformer_v2_N@12_L@6_M@2.yml │ │ │ └── equiformer_v2_N@12_L@6_M@2_epochs@30.yml │ │ └── all_md │ │ └── equiformer_v2 │ │ ├── equiformer_v2_N@20_L@6_M@3_153M.yml │ │ └── equiformer_v2_N@8_L@4_M@2_31M.yml └── trainer │ ├── __init__.py │ ├── base_trainer_oc20.py │ ├── base_trainer_v2.py │ ├── dist_setup.py │ ├── energy_trainer_v2.py │ ├── engine.py │ ├── forces_trainer_v2.py │ ├── lmdb_dataset.py │ ├── logger.py │ ├── lr_scheduler.py │ ├── make_lmdb_sizes.py │ ├── task_compute_stats.py │ └── task_relaxation.py ├── optim_factory.py ├── scripts └── train │ └── oc20 │ └── s2ef │ └── equiformer_v2 │ ├── equiformer_v2_N@12_L@6_M@2_splits@2M_g@8.sh │ ├── equiformer_v2_N@12_L@6_M@2_splits@2M_g@multi-nodes.sh │ ├── equiformer_v2_N@20_L@6_M@3_splits@all+md_g@multi-nodes.sh │ └── equiformer_v2_N@8_L@4_M@2_splits@all+md_g@multi-nodes.sh └── utils.py /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/README.md -------------------------------------------------------------------------------- /datasets/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /docs/changelog.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/docs/changelog.md -------------------------------------------------------------------------------- /docs/env_setup.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/docs/env_setup.md -------------------------------------------------------------------------------- /docs/equiformer_v2_poster.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/docs/equiformer_v2_poster.pdf -------------------------------------------------------------------------------- /engine.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/engine.py -------------------------------------------------------------------------------- /env/env_equiformer_v2.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/env/env_equiformer_v2.yml -------------------------------------------------------------------------------- /fig/equiformer_v2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2.png -------------------------------------------------------------------------------- /fig/equiformer_v2_adsorbml_results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_adsorbml_results.png -------------------------------------------------------------------------------- /fig/equiformer_v2_oc20_results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_oc20_results.png -------------------------------------------------------------------------------- /fig/equiformer_v2_oc22_results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_oc22_results.png -------------------------------------------------------------------------------- /fig/equiformer_v2_speed_accuracy_tradeoffs.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/fig/equiformer_v2_speed_accuracy_tradeoffs.png -------------------------------------------------------------------------------- /logger.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/logger.py -------------------------------------------------------------------------------- /main_oc20.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/main_oc20.py -------------------------------------------------------------------------------- /nets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/__init__.py -------------------------------------------------------------------------------- /nets/equiformer_v2/Jd.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/Jd.pt -------------------------------------------------------------------------------- /nets/equiformer_v2/activation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/activation.py -------------------------------------------------------------------------------- /nets/equiformer_v2/drop.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/drop.py -------------------------------------------------------------------------------- /nets/equiformer_v2/edge_rot_mat.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/edge_rot_mat.py -------------------------------------------------------------------------------- /nets/equiformer_v2/equiformer_v2_oc20.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/equiformer_v2_oc20.py -------------------------------------------------------------------------------- /nets/equiformer_v2/gaussian_rbf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/gaussian_rbf.py -------------------------------------------------------------------------------- /nets/equiformer_v2/input_block.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/input_block.py -------------------------------------------------------------------------------- /nets/equiformer_v2/layer_norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/layer_norm.py -------------------------------------------------------------------------------- /nets/equiformer_v2/module_list.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/module_list.py -------------------------------------------------------------------------------- /nets/equiformer_v2/radial_function.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/radial_function.py -------------------------------------------------------------------------------- /nets/equiformer_v2/so2_ops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/so2_ops.py -------------------------------------------------------------------------------- /nets/equiformer_v2/so3.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/so3.py -------------------------------------------------------------------------------- /nets/equiformer_v2/transformer_block.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/transformer_block.py -------------------------------------------------------------------------------- /nets/equiformer_v2/wigner.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/nets/equiformer_v2/wigner.py -------------------------------------------------------------------------------- /oc20/configs/s2ef/2M/base.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/configs/s2ef/2M/base.yml -------------------------------------------------------------------------------- /oc20/configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2.yml -------------------------------------------------------------------------------- /oc20/configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2_epochs@30.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2_epochs@30.yml -------------------------------------------------------------------------------- /oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@20_L@6_M@3_153M.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@20_L@6_M@3_153M.yml -------------------------------------------------------------------------------- /oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@8_L@4_M@2_31M.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@8_L@4_M@2_31M.yml -------------------------------------------------------------------------------- /oc20/trainer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/__init__.py -------------------------------------------------------------------------------- /oc20/trainer/base_trainer_oc20.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/base_trainer_oc20.py -------------------------------------------------------------------------------- /oc20/trainer/base_trainer_v2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/base_trainer_v2.py -------------------------------------------------------------------------------- /oc20/trainer/dist_setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/dist_setup.py -------------------------------------------------------------------------------- /oc20/trainer/energy_trainer_v2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/energy_trainer_v2.py -------------------------------------------------------------------------------- /oc20/trainer/engine.py: -------------------------------------------------------------------------------- 1 | ../../engine.py -------------------------------------------------------------------------------- /oc20/trainer/forces_trainer_v2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/forces_trainer_v2.py -------------------------------------------------------------------------------- /oc20/trainer/lmdb_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/lmdb_dataset.py -------------------------------------------------------------------------------- /oc20/trainer/logger.py: -------------------------------------------------------------------------------- 1 | ../../logger.py -------------------------------------------------------------------------------- /oc20/trainer/lr_scheduler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/lr_scheduler.py -------------------------------------------------------------------------------- /oc20/trainer/make_lmdb_sizes.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/make_lmdb_sizes.py -------------------------------------------------------------------------------- /oc20/trainer/task_compute_stats.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/task_compute_stats.py -------------------------------------------------------------------------------- /oc20/trainer/task_relaxation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/oc20/trainer/task_relaxation.py -------------------------------------------------------------------------------- /optim_factory.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/optim_factory.py -------------------------------------------------------------------------------- /scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@8.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@8.sh -------------------------------------------------------------------------------- /scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@multi-nodes.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@multi-nodes.sh -------------------------------------------------------------------------------- /scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@20_L@6_M@3_splits@all+md_g@multi-nodes.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@20_L@6_M@3_splits@all+md_g@multi-nodes.sh -------------------------------------------------------------------------------- /scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@8_L@4_M@2_splits@all+md_g@multi-nodes.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@8_L@4_M@2_splits@all+md_g@multi-nodes.sh -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/atomicarchitects/equiformer_v2/HEAD/utils.py --------------------------------------------------------------------------------