Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

omat24 example #295

Merged
merged 11 commits into from
Nov 9, 2024
Merged

omat24 example #295

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
195d5ce
omat24 example
allaffa Oct 21, 2024
38f6ebf
formatting fixed
allaffa Oct 22, 2024
a091e05
remove unused parsed arguments
allaffa Oct 23, 2024
8a2b477
exists set to True for data download
allaffa Oct 23, 2024
c54fc1b
exists set to True for data download
allaffa Oct 23, 2024
341ea73
updated download script to move tar files into subdirectories
Oct 25, 2024
7471611
bug fixed in normalization of the energy
allaffa Oct 26, 2024
cff9840
fixes on omat24 example
Nov 5, 2024
09f21dd
black formatting fixed
allaffa Nov 5, 2024
14c881b
Update examples/omat24/train.py
allaffa Nov 8, 2024
b1a7852
formatting fixed
allaffa Nov 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 85 additions & 0 deletions examples/omat24/download_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
import argparse
import glob
import logging
import os
import shutil


DOWNLOAD_LINKS = {
"train": {
"rattled-1000": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/rattled-1000.tar.gz",
"rattled-1000-subsampled": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/rattled-1000-subsampled.tar.gz",
"rattled-500": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/rattled-500.tar.gz",
"rattled-500-subsampled": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/rattled-500-subsampled.tar.gz",
"rattled-300": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/rattled-300.tar.gz",
"rattled-300-subsampled": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/rattled-300-subsampled.tar.gz",
"aimd-from-PBE-1000-npt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/aimd-from-PBE-1000-npt.tar.gz",
"aimd-from-PBE-1000-nvt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/aimd-from-PBE-1000-nvt.tar.gz",
"aimd-from-PBE-3000-npt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/aimd-from-PBE-3000-npt.tar.gz",
"aimd-from-PBE-3000-nvt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/aimd-from-PBE-3000-nvt.tar.gz",
"rattled-relax": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/train/rattled-relax.tar.gz",
},
"val": {
"rattled-1000": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/rattled-1000.tar.gz",
"rattled-1000-subsampled": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/rattled-1000-subsampled.tar.gz",
"rattled-500": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/rattled-500.tar.gz",
"rattled-500-subsampled": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/rattled-500-subsampled.tar.gz",
"rattled-300": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/rattled-300.tar.gz",
"rattled-300-subsampled": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/rattled-300-subsampled.tar.gz",
"aimd-from-PBE-1000-npt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/aimd-from-PBE-1000-npt.tar.gz",
"aimd-from-PBE-1000-nvt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/aimd-from-PBE-1000-nvt.tar.gz",
"aimd-from-PBE-3000-npt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/aimd-from-PBE-3000-npt.tar.gz",
"aimd-from-PBE-3000-nvt": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/aimd-from-PBE-3000-nvt.tar.gz",
"rattled-relax": "wget https://dl.fbaipublicfiles.com/opencatalystproject/data/omat/241018/omat/val/rattled-relax.tar.gz",
},
}


assert (
DOWNLOAD_LINKS["train"].keys() == DOWNLOAD_LINKS["val"].keys()
), "data partition names in train do not match with equivalent names in val"
dataset_names = list(DOWNLOAD_LINKS["train"].keys())


def get_data(datadir, task, split):
os.makedirs(datadir, exist_ok=True)

if (task == "train" or task == "val") and split is None:
raise NotImplementedError(f"{task} requires a split to be defined.")

assert (
split in DOWNLOAD_LINKS[task]
), f'{task}/{split}" split not defined, please specify one of the following: {list(DOWNLOAD_LINKS[task].keys())}'
download_link = DOWNLOAD_LINKS[task][split]

os.makedirs(os.path.join(datadir, task), exist_ok=True)

os.system(f"wget {download_link} -P {datadir}")
filename = os.path.join(datadir, os.path.basename(download_link))

# Move the directory
new_filename = os.path.join(datadir, task, os.path.basename(download_link))
shutil.move(filename, new_filename)

logging.info("Extracting contents...")
os.system(f"tar -xvf {new_filename} -C {os.path.join(datadir, task)}")


if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--data-path",
type=str,
default="./dataset",
help="Specify path to save datasets. Defaults to './dataset'",
)

args, _ = parser.parse_known_args()

for task in ["train", "val"]:
for split in dataset_names:
get_data(
datadir=args.data_path,
task=task,
split=split,
)
1 change: 1 addition & 0 deletions examples/omat24/example_requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
fairchem-core
58 changes: 58 additions & 0 deletions examples/omat24/omat24_energy.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
{
"Verbosity": {
"level": 2
},
"NeuralNetwork": {
"Architecture": {
"model_type": "EGNN",
"equivariance": true,
"radius": 5.0,
"max_neighbours": 100000,
"num_gaussians": 50,
"envelope_exponent": 5,
"int_emb_size": 64,
"basis_emb_size": 8,
"out_emb_size": 128,
"num_after_skip": 2,
"num_before_skip": 1,
"num_radial": 6,
"num_spherical": 7,
"num_filters": 126,
"edge_features": ["length"],
"hidden_dim": 50,
"num_conv_layers": 3,
"output_heads": {
"graph":{
"num_sharedlayers": 2,
"dim_sharedlayers": 50,
"num_headlayers": 2,
"dim_headlayers": [50,25]
}
},
"task_weights": [1.0]
},
"Variables_of_interest": {
"input_node_features": [0, 1, 2, 3],
"output_names": ["energy"],
"output_index": [0],
"output_dim": [1],
"type": ["graph"]
},
"Training": {
"num_epoch": 50,
"perc_train": 0.8,
"loss_function_type": "mae",
"batch_size": 32,
"continue": 0,
"Optimizer": {
"type": "AdamW",
"learning_rate": 1e-3
}
}
},
"Visualization": {
"plot_init_solution": true,
"plot_hist_solution": false,
"create_plots": true
}
}
58 changes: 58 additions & 0 deletions examples/omat24/omat24_forces.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
{
"Verbosity": {
"level": 2
},
"NeuralNetwork": {
"Architecture": {
"model_type": "EGNN",
"equivariance": true,
"radius": 5.0,
"max_neighbours": 100000,
"num_gaussians": 50,
"envelope_exponent": 5,
"int_emb_size": 64,
"basis_emb_size": 8,
"out_emb_size": 128,
"num_after_skip": 2,
"num_before_skip": 1,
"num_radial": 6,
"num_spherical": 7,
"num_filters": 126,
"edge_features": ["length"],
"hidden_dim": 50,
"num_conv_layers": 3,
"output_heads": {
"node": {
"num_headlayers": 2,
"dim_headlayers": [200,200],
"type": "mlp"
}
},
"task_weights": [1.0]
},
"Variables_of_interest": {
"input_node_features": [0, 1, 2, 3],
"output_names": ["forces"],
"output_index": [2],
"output_dim": [3],
"type": ["node"]
},
"Training": {
"num_epoch": 50,
"EarlyStopping": true,
"perc_train": 0.9,
"loss_function_type": "mae",
"batch_size": 32,
"continue": 0,
"Optimizer": {
"type": "AdamW",
"learning_rate": 1e-3
}
}
},
"Visualization": {
"plot_init_solution": true,
"plot_hist_solution": false,
"create_plots": true
}
}
Loading
Loading