├── .gitignore ├── LICENSE ├── README.md ├── assets └── HEGNN-teaser.jpg ├── datasets ├── __init__.py ├── md17 │ └── dataset.py └── nbody │ ├── __init__.py │ ├── datagen │ ├── generate_dataset.py │ ├── physical_objects.py │ ├── run.sh │ └── system.py │ ├── dataset.py │ └── script │ ├── generate_dataset_test.py │ ├── generate_dataset_train.py │ ├── generate_dataset_valid.py │ ├── physical_objects.py │ ├── run.sh │ ├── run_100_10.sh │ ├── run_100_3.sh │ └── system.py ├── expressivity ├── experiments │ ├── data │ │ └── reg_poly_coords.py │ └── utils │ │ ├── plot_utils.py │ │ └── train_utils.py ├── models │ ├── egnn.py │ ├── gvpgnn.py │ ├── hegnn.py │ ├── layers │ │ ├── egnn_layer.py │ │ ├── gvp_layer.py │ │ └── tfn_layer.py │ ├── mace.py │ ├── mace_modules │ │ ├── blocks.py │ │ ├── cg.py │ │ ├── irreps_tools.py │ │ ├── radial.py │ │ └── symmetric_contraction.py │ ├── schnet.py │ └── tfn.py ├── perturbation-test.ipynb ├── reg-poly-test.ipynb └── rot-3D-test.ipynb ├── main_md17.py ├── main_nbody.py ├── models ├── Clof │ ├── clof.py │ ├── gcl.py │ └── layers.py ├── DimeNet.py ├── EGNN.py ├── GVP.py ├── HEGNN.py ├── SchNet.py ├── __init__.py ├── basic.py ├── mace │ ├── blocks.py │ ├── cg.py │ ├── irreps_tools.py │ ├── mace.py │ ├── radial.py │ ├── symmetric_contraction.py │ └── tfn_layer.py ├── md17 │ ├── gcl.py │ └── layer.py ├── md17_model.py ├── se3_dynamics │ ├── __init__.py │ ├── dynamics.py │ ├── equivariant_attention │ │ ├── __init__.py │ │ ├── fibers.py │ │ ├── from_se3cnn │ │ │ ├── SO3.py │ │ │ ├── __init__.py │ │ │ ├── cache_file.py │ │ │ ├── license.txt │ │ │ ├── representations.py │ │ │ └── utils_steerable.py │ │ ├── modules.py │ │ └── ops.py │ ├── models.py │ └── utils │ │ ├── __init__.py │ │ └── utils_profiling.py └── segnn │ ├── balanced_irreps.py │ ├── instance_norm.py │ ├── o3_building_blocks.py │ └── segnn.py ├── requirements.txt └── utils ├── __init__.py ├── rotate.py ├── seed.py └── train.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/README.md -------------------------------------------------------------------------------- /assets/HEGNN-teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/assets/HEGNN-teaser.jpg -------------------------------------------------------------------------------- /datasets/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /datasets/md17/dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/md17/dataset.py -------------------------------------------------------------------------------- /datasets/nbody/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /datasets/nbody/datagen/generate_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/datagen/generate_dataset.py -------------------------------------------------------------------------------- /datasets/nbody/datagen/physical_objects.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/datagen/physical_objects.py -------------------------------------------------------------------------------- /datasets/nbody/datagen/run.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/datagen/run.sh -------------------------------------------------------------------------------- /datasets/nbody/datagen/system.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/datagen/system.py -------------------------------------------------------------------------------- /datasets/nbody/dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/dataset.py -------------------------------------------------------------------------------- /datasets/nbody/script/generate_dataset_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/generate_dataset_test.py -------------------------------------------------------------------------------- /datasets/nbody/script/generate_dataset_train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/generate_dataset_train.py -------------------------------------------------------------------------------- /datasets/nbody/script/generate_dataset_valid.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/generate_dataset_valid.py -------------------------------------------------------------------------------- /datasets/nbody/script/physical_objects.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/physical_objects.py -------------------------------------------------------------------------------- /datasets/nbody/script/run.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/run.sh -------------------------------------------------------------------------------- /datasets/nbody/script/run_100_10.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/run_100_10.sh -------------------------------------------------------------------------------- /datasets/nbody/script/run_100_3.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/run_100_3.sh -------------------------------------------------------------------------------- /datasets/nbody/script/system.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/datasets/nbody/script/system.py -------------------------------------------------------------------------------- /expressivity/experiments/data/reg_poly_coords.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/experiments/data/reg_poly_coords.py -------------------------------------------------------------------------------- /expressivity/experiments/utils/plot_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/experiments/utils/plot_utils.py -------------------------------------------------------------------------------- /expressivity/experiments/utils/train_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/experiments/utils/train_utils.py -------------------------------------------------------------------------------- /expressivity/models/egnn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/egnn.py -------------------------------------------------------------------------------- /expressivity/models/gvpgnn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/gvpgnn.py -------------------------------------------------------------------------------- /expressivity/models/hegnn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/hegnn.py -------------------------------------------------------------------------------- /expressivity/models/layers/egnn_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/layers/egnn_layer.py -------------------------------------------------------------------------------- /expressivity/models/layers/gvp_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/layers/gvp_layer.py -------------------------------------------------------------------------------- /expressivity/models/layers/tfn_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/layers/tfn_layer.py -------------------------------------------------------------------------------- /expressivity/models/mace.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/mace.py -------------------------------------------------------------------------------- /expressivity/models/mace_modules/blocks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/mace_modules/blocks.py -------------------------------------------------------------------------------- /expressivity/models/mace_modules/cg.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/mace_modules/cg.py -------------------------------------------------------------------------------- /expressivity/models/mace_modules/irreps_tools.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/mace_modules/irreps_tools.py -------------------------------------------------------------------------------- /expressivity/models/mace_modules/radial.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/mace_modules/radial.py -------------------------------------------------------------------------------- /expressivity/models/mace_modules/symmetric_contraction.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/mace_modules/symmetric_contraction.py -------------------------------------------------------------------------------- /expressivity/models/schnet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/schnet.py -------------------------------------------------------------------------------- /expressivity/models/tfn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/models/tfn.py -------------------------------------------------------------------------------- /expressivity/perturbation-test.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/perturbation-test.ipynb -------------------------------------------------------------------------------- /expressivity/reg-poly-test.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/reg-poly-test.ipynb -------------------------------------------------------------------------------- /expressivity/rot-3D-test.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/expressivity/rot-3D-test.ipynb -------------------------------------------------------------------------------- /main_md17.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/main_md17.py -------------------------------------------------------------------------------- /main_nbody.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/main_nbody.py -------------------------------------------------------------------------------- /models/Clof/clof.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/Clof/clof.py -------------------------------------------------------------------------------- /models/Clof/gcl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/Clof/gcl.py -------------------------------------------------------------------------------- /models/Clof/layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/Clof/layers.py -------------------------------------------------------------------------------- /models/DimeNet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/DimeNet.py -------------------------------------------------------------------------------- /models/EGNN.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/EGNN.py -------------------------------------------------------------------------------- /models/GVP.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/GVP.py -------------------------------------------------------------------------------- /models/HEGNN.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/HEGNN.py -------------------------------------------------------------------------------- /models/SchNet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/SchNet.py -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/basic.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/basic.py -------------------------------------------------------------------------------- /models/mace/blocks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/mace/blocks.py -------------------------------------------------------------------------------- /models/mace/cg.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/mace/cg.py -------------------------------------------------------------------------------- /models/mace/irreps_tools.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/mace/irreps_tools.py -------------------------------------------------------------------------------- /models/mace/mace.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/mace/mace.py -------------------------------------------------------------------------------- /models/mace/radial.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/mace/radial.py -------------------------------------------------------------------------------- /models/mace/symmetric_contraction.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/mace/symmetric_contraction.py -------------------------------------------------------------------------------- /models/mace/tfn_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/mace/tfn_layer.py -------------------------------------------------------------------------------- /models/md17/gcl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/md17/gcl.py -------------------------------------------------------------------------------- /models/md17/layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/md17/layer.py -------------------------------------------------------------------------------- /models/md17_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/md17_model.py -------------------------------------------------------------------------------- /models/se3_dynamics/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/se3_dynamics/dynamics.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/dynamics.py -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/fibers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/fibers.py -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/from_se3cnn/SO3.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/from_se3cnn/SO3.py -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/from_se3cnn/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/from_se3cnn/cache_file.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/from_se3cnn/cache_file.py -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/from_se3cnn/license.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/from_se3cnn/license.txt -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/from_se3cnn/representations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/from_se3cnn/representations.py -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/from_se3cnn/utils_steerable.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/from_se3cnn/utils_steerable.py -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/modules.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/modules.py -------------------------------------------------------------------------------- /models/se3_dynamics/equivariant_attention/ops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/equivariant_attention/ops.py -------------------------------------------------------------------------------- /models/se3_dynamics/models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/models.py -------------------------------------------------------------------------------- /models/se3_dynamics/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/se3_dynamics/utils/utils_profiling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/se3_dynamics/utils/utils_profiling.py -------------------------------------------------------------------------------- /models/segnn/balanced_irreps.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/segnn/balanced_irreps.py -------------------------------------------------------------------------------- /models/segnn/instance_norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/segnn/instance_norm.py -------------------------------------------------------------------------------- /models/segnn/o3_building_blocks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/segnn/o3_building_blocks.py -------------------------------------------------------------------------------- /models/segnn/segnn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/models/segnn/segnn.py -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/requirements.txt -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /utils/rotate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/utils/rotate.py -------------------------------------------------------------------------------- /utils/seed.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/utils/seed.py -------------------------------------------------------------------------------- /utils/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/GLAD-RUC/HEGNN/HEAD/utils/train.py --------------------------------------------------------------------------------