equiformer-v3 0.0.0__py3-none-any.whl

equiformer_v3/__init__.py

@@ -0,0 +1,6 @@
+# Register equiformer_v3 models and trainers with fairchem registry
+import equiformer_v3.experimental.models.equiformer_v3.equiformer_v3
+import equiformer_v3.experimental.models.equiformer_v3.equiformer_v3_dens
+import equiformer_v3.experimental.trainers.equiformer_v3_dens_trainer
+import equiformer_v3.experimental.trainers.dens_ase_dataset
+import equiformer_v3.experimental.trainers.oc20_total_energy_lmdb 

equiformer_v3/applications/AdsorbML/LICENSE.md

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) Facebook, Inc. and its affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

equiformer_v3/applications/AdsorbML/MODELS.md

@@ -0,0 +1,13 @@
+All pre-trained model checkpoints were obtained directly from [ocp](../core/model_checkpoints). For convenience, we provide a pointer to the checkpoints used in this work below:
+
+|Model |Download | ML+SP Success Rate @k=5 |
+| - | - | - |
+|SchNet |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2020_11/s2ef/schnet_all_large.pt) |5.04% |
+|DimeNet++ |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2021_02/s2ef/dimenetpp_all.pt) |10.79% |
+|PaiNN | [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_05/s2ef/painn_h512_s2ef_all.pt) |39.57% |
+|GemNet-OC | [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_07/s2ef/gemnet_oc_base_s2ef_all.pt) |82.94% |
+|GemNet-OC All+MD | [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/data/gemnet_oc_s2ef_all_md.pt) |84.38% |
+|GemNet-OC-Large All+MD | [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2022_07/s2ef/gemnet_oc_large_s2ef_all_md.pt) |86.02% |
+|SCN-Large All+MD | [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/data/scn_all_md_s2ef.pt) |87.77% |
+
+For more details please visit ../core/model_checkpoints.

equiformer_v3/applications/AdsorbML/README.md

@@ -0,0 +1,65 @@
+## AdsorbML: Accelerating Adsorption Energy Calculations with Machine Learning
+
+![adsorbml](https://user-images.githubusercontent.com/45150244/213025581-498b459e-9077-42ac-84e1-65ef331555d2.png)
+
+`AdsorbML` is an algorithm to calculating the minima adsorbate binding energy (adsorption energy) for a unique adsorbate+surface combination. All ML models are obtained from [`ocp`](https://github.com/Open-Catalyst-Project/ocp) to perform corresponding structure relaxations.
+
+This repository holds the dataset, scripts, and downloads for the accompanying [paper](https://arxiv.org/abs/2211.16486).
+
+### OC20-Dense Dataset (OC20-Dense)
+
+OC20-Dense contains a dense sampling of adsorbate configurations on ~1,000 randomly selected adsorbate+surface materials from the [OC20](https://arxiv.org/abs/2010.09990) dataset. The dataset comprises a total of 85,658 unique input configurations.
+
+The dataset is stored in an LMDB file and ready to be used in `ocp` upon download. Additionally, ground truth DFT relaxations are store in ASE trajectories and provided for all converged systems used for evaluation.
+
+NOTE - ASE trajectories exclude systems that were not converged or had invalid configurations as defined by the constraints in the `AdsorbML` manuscript. This resulted in 65,073 relaxations available for evaluation and are provided here.
+
+
+|Splits |Size of compressed version (in bytes)  |Size of uncompressed version (in bytes)    | MD5 checksum (download link)   |
+|---    |---    |---    |---    |
+|LMDB    |654M   |9.8G   | [0163b0e8c4df6d9c426b875a28d9178a](https://dl.fbaipublicfiles.com/opencatalystproject/data/adsorbml/oc20_dense_data.tar.gz)   |
+|ASE Trajectories    |29G    |112G   | [ee937e5290f8f720c914dc9a56e0281f](https://dl.fbaipublicfiles.com/opencatalystproject/data/adsorbml/oc20_dense_trajectories.tar.gz)   |
+
+The following files are also provided to be used for evaluation and general information:
+* `oc20dense_mapping.pkl` : Mapping of the LMDB `sid` to general metadata information. If this file is not present, run the command `python src/fairchem/core/scripts/download_large_files.py adsorbml` from the root of the fairchem repo to download it. -
+  * `system_id`: Unique system identifier for an adsorbate, bulk, surface combination.
+  * `config_id`: Unique configuration identifier, where `rand` and `heur` correspond to random and heuristic initial configurations, respectively.
+  * `mpid`: Materials Project bulk identifier.
+  * `miller_idx`: 3-tuple of integers indicating the Miller indices of the surface.
+  * `shift`: C-direction shift used to determine cutoff for the surface (c-direction is following the nomenclature from Pymatgen).
+  * `top`: Boolean indicating whether the chosen surface was at the top or bottom of the originally enumerated surface.
+  * `adsorbate`: Chemical composition of the adsorbate.
+  * `adsorption_site`: A tuple of 3-tuples containing the Cartesian coordinates of each binding adsorbate atom
+* `oc20dense_targets.pkl` :  DFT adsorption energies across different system and placement ids.
+* `oc20dense_compute.pkl` :  DFT compute as measured in the number of ionic and scf steps for each evaluated relaxation.
+* `oc20dense_ref_energies.pkl` : Reference energy used for a specified `system_id`. This energy includes the relaxed clean surface and the gas phase adsorbate energy to ensure consistency across calculations.
+* `oc20dense_tags.pkl` : Tag information used for a specified `system_id`. Where 0 = subsurface, 1 = surface, 2 = adsorbate.
+
+All mappings can be obtained at the following downloadable link: https://dl.fbaipublicfiles.com/opencatalystproject/data/adsorbml/oc20_dense_mappings.tar.gz
+
+MD5 checksums:
+```
+c18735c405ce6ce5761432b07287d8d9  oc20_dense_mappings.tar.gz
+3e26c3bcef01ccfc9b001931065ea6e6  oc20dense_mapping.pkl
+fd589b013b72e62e11a6b2a5bd1d323c  oc20dense_targets.pkl
+78d25997e0aaf754df526ab37276bb89  oc20dense_compute.pkl
+b07c64158e4bfa5f7b9bf6263753ecc5  oc20dense_ref_energies.pkl
+1ba0bc266130f186850f5faa547b6a02  oc20dense_tags.pkl
+```
+
+### Running `AdsorbML`
+
+Please see the [README](adsorbml/scripts/README.md) inside the `scripts` directory for instructions.
+
+### Citing `AdsorbML`
+
+If you use this codebase in your work, please consider citing:
+
+```bibtex
+@article{lan2022adsorbml,
+  title={AdsorbML: Accelerating Adsorption Energy Calculations with Machine Learning},
+  author={Lan*, Janice and Palizhati*, Aini and Shuaibi*, Muhammed and Wood*, Brandon M and Wander, Brook and Das, Abhishek and Uyttendaele, Matt and Zitnick, C Lawrence and Ulissi, Zachary W},
+  journal={arXiv preprint arXiv:2211.16486},
+  year={2022}
+}
+```

equiformer_v3/applications/AdsorbML/adsorbml/2023_neurips_challenge/README.md

@@ -0,0 +1,17 @@
+## Validating energy predictions for 2023 Open Catalyst Challenge
+
+The `challenge_eval.py` script takes in your prediction npz file and the model used to generate the ML relaxed strcutures (gemnet-oc-2M, scn-2M, or escn-2M) and returns the success rate. More details on how to run energy predictions on ML relaxed strcutures can be found on the [challenge website](https://opencatalystproject.org/challenge.html) under the evaluation section.
+
+1. Git clone this repository:
+    ```
+    git clone https://github.com/Open-Catalyst-Project/AdsorbML.git
+    ```
+2. Change into the 2023_neurips_challenge directory:
+    ```
+    cd AdsorbML/adsorbml/2023_neurips_challenge
+    ```
+3. Run script:
+    ```
+    python challenge_eval.py --model model_used_for_MLRS --results-file /path/to/predictions.npz
+    ```
+    The `--model` variable should be set to either `gemnet-oc-2M`, `scn-2M`, or `escn-2M` depending on which LMDB you chose.

equiformer_v3/applications/AdsorbML/adsorbml/2023_neurips_challenge/challenge_eval.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import argparse
+import pickle
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+
+from fairchem.core.scripts import download_large_files
+
+
+def is_successful(best_pred_energy, best_dft_energy, SUCCESS_THRESHOLD=0.1):
+    """
+    Computes the success rate given the best predicted energy
+    and the best ground truth DFT energy.
+
+    success_parity: The standard definition for success, where ML needs to be
+    within the SUCCESS_THRESHOLD, or lower, of the DFT energy.
+
+    Returns: Bool
+    """
+    # Given best ML and DFT energy, compute various success metrics:
+    # success_parity: base success metric (ML - DFT <= SUCCESS_THRESHOLD)
+    diff = best_pred_energy - best_dft_energy
+    return diff <= SUCCESS_THRESHOLD
+
+
+def compute_valid_ml_success(ml_data, dft_data):
+    """
+    Computes validated ML success rates.
+    Here, results are generated only from ML. DFT single-points are used to
+    validate whether the ML energy is within 0.1eV of the DFT energy of the
+    predicted structure. If valid, the ML energy is compared to the ground
+    truth DFT energy, otherwise it is discarded.
+
+    Return validated ML success rates.
+    """
+
+    success_rate = 0.0
+
+    for system in dft_data:
+        # For `system`, collect all ML adslabs and their corresponding energies
+        ml_adslabs, ml_energies = [], []
+        for config in ml_data[system]:
+            ml_adslabs.append(config)
+            ml_energies.append(ml_data[system][config]["ml_energy"])
+
+        min_ml_idx = np.argmin(ml_energies)
+        min_adslab = ml_adslabs[min_ml_idx]
+        best_ml_energy = ml_energies[min_ml_idx]
+        # If the best ML energy is not within 0.1eV
+        # of its DFT energy evaluation, discard.
+        ml_dft_energy = ml_data[system][min_adslab]["ml+dft_energy"]
+        diff = abs(ml_dft_energy - best_ml_energy)
+        if diff > 0.1:
+            continue
+
+        best_dft_energy = min(list(dft_data[system].values()))
+
+        success = is_successful(best_ml_energy, best_dft_energy)
+        success_rate += success
+
+    success_rate /= len(dft_data)
+
+    print("=" * 50)
+    print(f"Success Rate (%): {100*success_rate}")
+
+
+def get_dft_data(targets):
+    """
+    Organizes the released target mapping for evaluation lookup.
+
+    Returns: Dict:
+        {
+           'system_id 1': {'config_id 1': dft_ads_energy, 'config_id 2': dft_ads_energy},
+           'system_id 2': {'config_id 1': dft_ads_energy, 'config_id 2': dft_ads_energy},
+           ...
+        }
+    """
+    dft_data = defaultdict(dict)
+    for system in targets:
+        for adslab in targets[system]:
+            dft_data[system][adslab[0]] = adslab[1]
+
+    return dft_data
+
+
+def process_ml_data(results_file, model, metadata, ml_dft_targets, dft_data):
+    """
+    For ML systems in which no configurations made it through the physical
+    constraint checks, set energies to an arbitrarily high value to ensure
+    a failure case in evaluation.
+
+    Returns: Dict:
+        {
+           'system_id 1': {'config_id 1': {'ml_energy': predicted energy, 'ml+dft_energy': dft energy of ML structure} ...},
+           'system_id 2': {'config_id 1': {'ml_energy': predicted energy, 'ml+dft_energy': dft energy of ML structure} ...},
+           ...
+        }
+    """
+    preds = np.load(results_file)
+    ml_data = defaultdict(dict)
+
+    for _id, energy in zip(preds["ids"], preds["energy"]):
+        sid, _ = _id.split("_")
+
+        info = metadata[int(sid)]
+        sysid = info["system_id"]
+        config = info["config_id"]
+
+        ml_dft_energy = ml_dft_targets[model][sysid][config]
+        ml_data[sysid][config] = {"ml_energy": energy, "ml+dft_energy": ml_dft_energy}
+
+    # set missing systems to high energy
+    # set missing systems to 0 DFT compute
+    for system in dft_data:
+        if system not in ml_data:
+            ml_data[system] = defaultdict(dict)
+        for config in dft_data[system]:
+            if config not in ml_data[system]:
+                _dict = {
+                    "ml_energy": 1e10,
+                    "ml+dft_energy": 1e10,
+                }
+                ml_data[system][config] = _dict
+
+    # for ML systems with no available ml+dft datapoints, set to an arbitrarily
+    # high energy value
+    for system in ml_data:
+        for config in ml_data[system]:
+            if not ml_data[system][config]["ml+dft_energy"]:
+                ml_data[system][config]["ml+dft_energy"] = 1e10
+
+    return ml_data
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        required=True,
+        choices=["gemnet-oc-2M", "escn-2M", "scn-2M"],
+    )
+    parser.add_argument(
+        "--results-file",
+        required=True,
+        help="Path to predictions to evaluate. NPZ format.",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    """
+    This script takes in your prediction file (npz format)
+    and the ML model name used for ML relaxations.
+    Then using a mapping file, dft ground truth energy,
+    and ML relaxed dft energy returns the success rate of your predictions.
+    """
+
+    args = parse_args()
+
+    # targets and metadata are expected to be in
+    # the same directory as this script
+    if (
+        not Path(__file__).with_name("oc20dense_val_targets.pkl").exists()
+        or not Path(__file__).with_name("ml_relaxed_dft_targets.pkl").exists()
+    ):
+        download_large_files.download_file_group("adsorbml")
+    targets = pickle.load(
+        open(Path(__file__).with_name("oc20dense_val_targets.pkl"), "rb")
+    )
+    ml_dft_targets = pickle.load(
+        open(Path(__file__).with_name("ml_relaxed_dft_targets.pkl"), "rb")
+    )
+    metadata = pickle.load(
+        open(Path(__file__).with_name("oc20dense_mapping.pkl"), "rb")
+    )
+
+    ###### Process DFT Data ######
+    dft_data = get_dft_data(targets)
+
+    ###### Process ML Data ######
+    ml_data = process_ml_data(
+        args.results_file, args.model, metadata, ml_dft_targets, dft_data
+    )
+
+    ###### Compute Metrics ######
+    print(f"Prediction file: {args.results_file}")
+    compute_valid_ml_success(ml_data, dft_data)
+
+
+if __name__ == "__main__":
+    main()

equiformer_v3/applications/AdsorbML/adsorbml/configs/dpp.yml

@@ -0,0 +1,61 @@
+trainer: forces
+
+dataset:
+  - src: data/s2ef/all/train/
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  - src: data/s2ef/all/val_id_30k/
+
+logger: wandb
+
+task:
+  dataset: trajectory_lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  eval_relaxations: True
+  relaxation_steps: 300
+  relaxation_fmax: 0.02
+  write_pos: False
+  relax_dataset:
+    src: path/to/oc20-dense/dataset #TODO
+  relax_opt:
+    name: lbfgs
+    maxstep: 0.04
+    memory: 50
+    damping: 1.0
+    alpha: 70.0
+    traj_dir: path/to/save/ml/relaxations #TODO
+
+model:
+  name: dimenetplusplus
+  hidden_channels: 192
+  out_emb_channels: 192
+  num_blocks: 3
+  cutoff: 6.0
+  num_radial: 6
+  num_spherical: 7
+  num_before_skip: 1
+  num_after_skip: 2
+  num_output_layers: 3
+  regress_forces: True
+  use_pbc: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  eval_every: 10000
+  num_workers: 8
+  lr_initial: 0.0001
+  lr_gamma: 0.1
+  lr_milestones: # steps at which lr_initial <- lr_initial * lr_gamma
+    - 130794
+    - 196192
+    - 261589
+  warmup_steps: 130794
+  warmup_factor: 0.2
+  max_epochs: 7
+  force_coefficient: 50

equiformer_v3/applications/AdsorbML/adsorbml/configs/gemnet-oc-large.yml

@@ -0,0 +1,109 @@
+trainer: forces
+
+dataset:
+  - src: data/s2ef/all/train/
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  - src: data/s2ef/all/val_id_30k/
+
+logger: wandb
+
+task:
+  dataset: trajectory_lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  eval_relaxations: True
+  relaxation_steps: 300
+  relaxation_fmax: 0.02
+  write_pos: False
+  relax_dataset:
+    src: path/to/oc20-dense/dataset #TODO
+  relax_opt:
+    name: lbfgs
+    maxstep: 0.04
+    memory: 50
+    damping: 1.0
+    alpha: 70.0
+    traj_dir: path/to/save/ml/relaxations #TODO
+
+model:
+  name: gemnet_oc
+  num_spherical: 7
+  num_radial: 128
+  num_blocks: 6
+  emb_size_atom: 256
+  emb_size_edge: 1024
+  emb_size_trip_in: 64
+  emb_size_trip_out: 128
+  emb_size_quad_in: 64
+  emb_size_quad_out: 32
+  emb_size_aint_in: 64
+  emb_size_aint_out: 64
+  emb_size_rbf: 32
+  emb_size_cbf: 16
+  emb_size_sbf: 64
+  num_before_skip: 2
+  num_after_skip: 2
+  num_concat: 4
+  num_atom: 3
+  num_output_afteratom: 3
+  cutoff: 12.0
+  cutoff_qint: 12.0
+  cutoff_aeaint: 12.0
+  cutoff_aint: 12.0
+  max_neighbors: 30
+  max_neighbors_qint: 8
+  max_neighbors_aeaint: 20
+  max_neighbors_aint: 1000
+  rbf:
+    name: gaussian
+  envelope:
+    name: polynomial
+    exponent: 5
+  cbf:
+    name: spherical_harmonics
+  sbf:
+    name: legendre_outer
+  extensive: True
+  output_init: HeOrthogonal
+  activation: silu
+  scale_file: configs/s2ef/all/gemnet/scaling_factors/gemnet-oc-large.pt
+
+  regress_forces: True
+  direct_forces: True
+  forces_coupled: False
+
+  quad_interaction: True
+  atom_edge_interaction: True
+  edge_atom_interaction: True
+  atom_interaction: True
+
+  num_atom_emb_layers: 2
+  num_global_out_layers: 2
+  qint_tags: [1, 2]
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 5000
+  num_workers: 2
+  lr_initial: 2.e-4
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  patience: 3
+  max_epochs: 80
+  force_coefficient: 100
+  energy_coefficient: 1
+  ema_decay: 0.999
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  weight_decay: 0

equiformer_v3/applications/AdsorbML/adsorbml/configs/gemnet-oc.yml

@@ -0,0 +1,109 @@
+trainer: forces
+
+dataset:
+  - src: data/s2ef/all/train/
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  - src: data/s2ef/all/val_id_30k/
+
+logger: wandb
+
+task:
+  dataset: trajectory_lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  eval_relaxations: True
+  relaxation_steps: 300
+  relaxation_fmax: 0.02
+  write_pos: False
+  relax_dataset:
+    src: path/to/oc20-dense/dataset #TODO
+  relax_opt:
+    name: lbfgs
+    maxstep: 0.04
+    memory: 50
+    damping: 1.0
+    alpha: 70.0
+    traj_dir: path/to/save/ml/relaxations #TODO
+
+model:
+  name: gemnet_oc
+  num_spherical: 7
+  num_radial: 128
+  num_blocks: 4
+  emb_size_atom: 256
+  emb_size_edge: 512
+  emb_size_trip_in: 64
+  emb_size_trip_out: 64
+  emb_size_quad_in: 32
+  emb_size_quad_out: 32
+  emb_size_aint_in: 64
+  emb_size_aint_out: 64
+  emb_size_rbf: 16
+  emb_size_cbf: 16
+  emb_size_sbf: 32
+  num_before_skip: 2
+  num_after_skip: 2
+  num_concat: 1
+  num_atom: 3
+  num_output_afteratom: 3
+  cutoff: 12.0
+  cutoff_qint: 12.0
+  cutoff_aeaint: 12.0
+  cutoff_aint: 12.0
+  max_neighbors: 30
+  max_neighbors_qint: 8
+  max_neighbors_aeaint: 20
+  max_neighbors_aint: 1000
+  rbf:
+    name: gaussian
+  envelope:
+    name: polynomial
+    exponent: 5
+  cbf:
+    name: spherical_harmonics
+  sbf:
+    name: legendre_outer
+  extensive: True
+  output_init: HeOrthogonal
+  activation: silu
+  scale_file: configs/s2ef/all/gemnet/scaling_factors/gemnet-oc.pt
+
+  regress_forces: True
+  direct_forces: True
+  forces_coupled: False
+
+  quad_interaction: True
+  atom_edge_interaction: True
+  edge_atom_interaction: True
+  atom_interaction: True
+
+  num_atom_emb_layers: 2
+  num_global_out_layers: 2
+  qint_tags: [1, 2]
+
+optim:
+  batch_size: 16
+  eval_batch_size: 16
+  load_balancing: atoms
+  eval_every: 5000
+  num_workers: 2
+  lr_initial: 5.e-4
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  patience: 3
+  max_epochs: 80
+  force_coefficient: 100
+  energy_coefficient: 1
+  ema_decay: 0.999
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  weight_decay: 0

equiformer_v3/applications/AdsorbML/adsorbml/configs/painn.yml

@@ -0,0 +1,66 @@
+trainer: forces
+
+dataset:
+  - src: data/s2ef/all/train/
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  - src: data/s2ef/all/val_id_30k/
+
+logger: wandb
+
+task:
+  dataset: trajectory_lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  eval_relaxations: True
+  relaxation_steps: 300
+  relaxation_fmax: 0.02
+  write_pos: False
+  relax_dataset:
+    src: path/to/oc20-dense/dataset #TODO
+  relax_opt:
+    name: lbfgs
+    maxstep: 0.04
+    memory: 50
+    damping: 1.0
+    alpha: 70.0
+    traj_dir: path/to/save/ml/relaxations #TODO
+
+model:
+  name: painn
+  hidden_channels: 512
+  num_layers: 6
+  num_rbf: 128
+  cutoff: 12.0
+  max_neighbors: 50
+  scale_file: configs/s2ef/all/painn/painn_nb6_scaling_factors.pt
+  regress_forces: True
+  direct_forces: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 5000
+  num_workers: 2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  lr_initial: 1.e-4
+  lr_gamma: 0.8
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  patience: 3
+  max_epochs: 80
+  force_coefficient: 100
+  energy_coefficient: 1
+  ema_decay: 0.999
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  weight_decay: 0  # 2e-6 (TF weight decay) / 1e-4 (lr) = 2e-2