AdsorbML tutorial#
from fairchem.core.common.relaxation.ase_utils import OCPCalculator
import ase.io
from ase.optimize import BFGS
from fairchem.data.oc.core import Adsorbate, AdsorbateSlabConfig, Bulk, Slab
import os
from glob import glob
import pandas as pd
from fairchem.data.oc.utils import DetectTrajAnomaly
from fairchem.data.oc.utils.vasp import write_vasp_input_files
# Optional - see below
import numpy as np
from dscribe.descriptors import SOAP
from scipy.spatial.distance import pdist, squareform
from x3dase.visualize import view_x3d_n
Enumerate the adsorbate-slab configurations to run relaxations on#
AdsorbML incorporates random placement, which is especially useful for more complicated adsorbates which may have many degrees of freedom. I have opted sample a few random placements and a few heuristic. Here I am using *CO on copper (1,1,1) as an example.
bulk_src_id = "mp-30"
adsorbate_smiles = "*CO"
bulk = Bulk(bulk_src_id_from_db = bulk_src_id)
adsorbate = Adsorbate(adsorbate_smiles_from_db=adsorbate_smiles)
slabs = Slab.from_bulk_get_specific_millers(bulk = bulk, specific_millers=(1,1,1))
# There may be multiple slabs with this miller index.
# For demonstrative purposes we will take the first entry.
slab = slabs[0]
Downloading src/fairchem/data/oc/databases/pkls/bulks.pkl...
# Perform heuristic placements
heuristic_adslabs = AdsorbateSlabConfig(slabs[0], adsorbate, mode="heuristic")
# Perform random placements
# (for AdsorbML we use `num_sites = 100` but we will use 4 for brevity here)
random_adslabs = AdsorbateSlabConfig(slabs[0], adsorbate, mode="random_site_heuristic_placement", num_sites = 4)
Run ML relaxations:#
There are 2 options for how to do this.
Using
OCPCalculatoras the calculator within the ASE frameworkBy writing objects to lmdb and relaxing them using
main.pyin the ocp repo
(1) is really only adequate for small stuff and it is what I will show here, but if you plan to run many relaxations, you should definitely use (2). More details about writing lmdbs has been provided here - follow the IS2RS/IS2RE instructions. And more information about running relaxations once the lmdb has been written is here.
You need to provide the calculator with a path to a model checkpoint file. That can be downloaded here
from fairchem.core.common.relaxation.ase_utils import OCPCalculator
from fairchem.core.models.model_registry import model_name_to_local_file
import os
checkpoint_path = model_name_to_local_file('EquiformerV2-31M-S2EF-OC20-All+MD', local_cache='/tmp/fairchem_checkpoints/')
os.makedirs(f"data/{bulk}_{adsorbate}", exist_ok=True)
# Define the calculator
calc = OCPCalculator(checkpoint_path=checkpoint_path) # if you have a gpu, add `cpu=False` to speed up calculations
adslabs = [*heuristic_adslabs.atoms_list, *random_adslabs.atoms_list]
# Set up the calculator
for idx, adslab in enumerate(adslabs):
adslab.calc = calc
opt = BFGS(adslab, trajectory=f"data/{bulk}_{adsorbate}/{idx}.traj")
opt.run(fmax=0.05, steps=100) # For the AdsorbML results we used fmax = 0.02 and steps = 300, but we will use less strict values for brevity.
/home/runner/work/fairchem/fairchem/src/fairchem/core/models/escn/so3.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_Jd = torch.load(os.path.join(os.path.dirname(__file__), "Jd.pt"))
/home/runner/work/fairchem/fairchem/src/fairchem/core/models/scn/spherical_harmonics.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_Jd = torch.load(os.path.join(os.path.dirname(__file__), "Jd.pt"))
/home/runner/work/fairchem/fairchem/src/fairchem/core/models/equiformer_v2/wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_Jd = torch.load(os.path.join(os.path.dirname(__file__), "Jd.pt"))
/home/runner/work/fairchem/fairchem/src/fairchem/core/models/equiformer_v2/layer_norm.py:75: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
@torch.cuda.amp.autocast(enabled=False)
/home/runner/work/fairchem/fairchem/src/fairchem/core/models/equiformer_v2/layer_norm.py:175: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
@torch.cuda.amp.autocast(enabled=False)
/home/runner/work/fairchem/fairchem/src/fairchem/core/models/equiformer_v2/layer_norm.py:263: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
@torch.cuda.amp.autocast(enabled=False)
/home/runner/work/fairchem/fairchem/src/fairchem/core/models/equiformer_v2/layer_norm.py:357: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
@torch.cuda.amp.autocast(enabled=False)
/home/runner/work/fairchem/fairchem/src/fairchem/core/common/relaxation/ase_utils.py:150: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
WARNING:root:Detected old config, converting to new format. Consider updating to avoid potential incompatibilities.
INFO:root:amp: true
cmd:
checkpoint_dir: /home/runner/work/fairchem/fairchem/docs/tutorials/checkpoints/2024-11-19-06-53-52
commit: aa298ac
identifier: ''
logs_dir: /home/runner/work/fairchem/fairchem/docs/tutorials/logs/wandb/2024-11-19-06-53-52
print_every: 100
results_dir: /home/runner/work/fairchem/fairchem/docs/tutorials/results/2024-11-19-06-53-52
seed: null
timestamp_id: 2024-11-19-06-53-52
version: 0.1.dev1+gaa298ac
dataset:
format: trajectory_lmdb_v2
grad_target_mean: 0.0
grad_target_std: 2.887317180633545
key_mapping:
force: forces
y: energy
normalize_labels: true
target_mean: -0.7554450631141663
target_std: 2.887317180633545
transforms:
normalizer:
energy:
mean: -0.7554450631141663
stdev: 2.887317180633545
forces:
mean: 0.0
stdev: 2.887317180633545
evaluation_metrics:
metrics:
energy:
- mae
forces:
- forcesx_mae
- forcesy_mae
- forcesz_mae
- mae
- cosine_similarity
- magnitude_error
misc:
- energy_forces_within_threshold
primary_metric: forces_mae
gp_gpus: null
gpus: 0
logger: wandb
loss_functions:
- energy:
coefficient: 4
fn: mae
- forces:
coefficient: 100
fn: l2mae
model:
alpha_drop: 0.1
attn_activation: silu
attn_alpha_channels: 64
attn_hidden_channels: 64
attn_value_channels: 16
distance_function: gaussian
drop_path_rate: 0.1
edge_channels: 128
ffn_activation: silu
ffn_hidden_channels: 128
grid_resolution: 18
lmax_list:
- 4
max_neighbors: 20
max_num_elements: 90
max_radius: 12.0
mmax_list:
- 2
name: equiformer_v2
norm_type: layer_norm_sh
num_distance_basis: 512
num_heads: 8
num_layers: 8
num_sphere_samples: 128
otf_graph: true
proj_drop: 0.0
regress_forces: true
sphere_channels: 128
use_atom_edge_embedding: true
use_gate_act: false
use_grid_mlp: true
use_pbc: true
use_s2_act_attn: false
weight_init: uniform
optim:
batch_size: 8
clip_grad_norm: 100
ema_decay: 0.999
energy_coefficient: 4
eval_batch_size: 8
eval_every: 10000
force_coefficient: 100
grad_accumulation_steps: 1
load_balancing: atoms
loss_energy: mae
loss_force: l2mae
lr_initial: 0.0004
max_epochs: 3
num_workers: 8
optimizer: AdamW
optimizer_params:
weight_decay: 0.001
scheduler: LambdaLR
scheduler_params:
epochs: 1009275
lambda_type: cosine
lr: 0.0004
lr_min_factor: 0.01
warmup_epochs: 3364.25
warmup_factor: 0.2
outputs:
energy:
level: system
forces:
eval_on_free_atoms: true
level: atom
train_on_free_atoms: true
relax_dataset: {}
slurm:
additional_parameters:
constraint: volta32gb
cpus_per_task: 9
folder: /checkpoint/abhshkdz/open-catalyst-project/logs/equiformer_v2/8307793
gpus_per_node: 8
job_id: '8307793'
job_name: eq2s_051701_allmd
mem: 480GB
nodes: 8
ntasks_per_node: 8
partition: learnaccel
time: 4320
task:
dataset: trajectory_lmdb_v2
eval_on_free_atoms: true
grad_input: atomic forces
labels:
- potential energy
primary_metric: forces_mae
train_on_free_atoms: true
test_dataset: {}
trainer: ocp
val_dataset: {}
INFO:root:Loading model: equiformer_v2
WARNING:root:equiformer_v2 (EquiformerV2) class is deprecated in favor of equiformer_v2_backbone_and_heads (EquiformerV2BackboneAndHeads)
INFO:root:Loaded EquiformerV2 with 31058690 parameters.
INFO:root:Loading checkpoint in inference-only mode, not loading keys associated with trainer state!
/home/runner/work/fairchem/fairchem/src/fairchem/core/modules/normalization/normalizer.py:69: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
"mean": torch.tensor(state_dict["mean"]),
WARNING:root:No seed has been set in modelcheckpoint or OCPCalculator! Results may not be reproducible on re-run
/home/runner/work/fairchem/fairchem/src/fairchem/core/trainers/ocp_trainer.py:461: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
with torch.cuda.amp.autocast(enabled=self.scaler is not None):
Step Time Energy fmax
BFGS: 0 06:54:23 -0.372193 1.894494
BFGS: 1 06:54:26 -0.375064 2.304425
BFGS: 2 06:54:29 -0.410459 0.838395
BFGS: 3 06:54:31 -0.439837 0.833875
BFGS: 4 06:54:34 -0.497138 1.265293
BFGS: 5 06:54:36 -0.537742 0.227679
BFGS: 6 06:54:39 -0.537880 0.188398
BFGS: 7 06:54:41 -0.539653 0.256894
BFGS: 8 06:54:44 -0.543899 0.347834
BFGS: 9 06:54:47 -0.551531 0.376806
BFGS: 10 06:54:49 -0.555696 0.253157
BFGS: 11 06:54:52 -0.558139 0.114778
BFGS: 12 06:54:54 -0.558915 0.154769
BFGS: 13 06:54:57 -0.559273 0.208317
BFGS: 14 06:55:00 -0.560790 0.217677
BFGS: 15 06:55:02 -0.560815 0.183834
BFGS: 16 06:55:05 -0.562498 0.079821
BFGS: 17 06:55:08 -0.563292 0.061306
BFGS: 18 06:55:10 -0.563837 0.108345
BFGS: 19 06:55:13 -0.563807 0.118470
BFGS: 20 06:55:16 -0.565091 0.096864
BFGS: 21 06:55:18 -0.565422 0.051761
BFGS: 22 06:55:21 -0.565065 0.048092
Step Time Energy fmax
BFGS: 0 06:55:23 -0.370925 1.325423
BFGS: 1 06:55:26 -0.378140 1.677682
BFGS: 2 06:55:29 -0.408270 1.049692
BFGS: 3 06:55:31 -0.506154 1.098950
BFGS: 4 06:55:34 -0.511561 0.363479
BFGS: 5 06:55:37 -0.511565 0.160638
BFGS: 6 06:55:39 -0.511986 0.325812
BFGS: 7 06:55:42 -0.513830 0.396905
BFGS: 8 06:55:44 -0.516858 0.259304
BFGS: 9 06:55:47 -0.517177 0.091769
BFGS: 10 06:55:50 -0.518391 0.133440
BFGS: 11 06:55:52 -0.518709 0.309994
BFGS: 12 06:55:55 -0.520680 0.369503
BFGS: 13 06:55:58 -0.522911 0.243706
BFGS: 14 06:56:00 -0.523245 0.073292
BFGS: 15 06:56:03 -0.523903 0.038656
Step Time Energy fmax
BFGS: 0 06:56:05 -0.225231 1.071860
BFGS: 1 06:56:08 -0.256597 1.505636
BFGS: 2 06:56:11 -0.350407 1.787051
BFGS: 3 06:56:13 -0.434558 1.486797
BFGS: 4 06:56:16 -0.468801 0.703190
BFGS: 5 06:56:18 -0.487961 0.327919
BFGS: 6 06:56:21 -0.495467 0.219621
BFGS: 7 06:56:24 -0.497261 0.207918
BFGS: 8 06:56:26 -0.513515 0.352554
BFGS: 9 06:56:29 -0.516381 0.204757
BFGS: 10 06:56:31 -0.518768 0.113466
BFGS: 11 06:56:34 -0.520004 0.196055
BFGS: 12 06:56:37 -0.522778 0.287858
BFGS: 13 06:56:39 -0.525264 0.315524
BFGS: 14 06:56:42 -0.527444 0.254864
BFGS: 15 06:56:44 -0.528471 0.180526
BFGS: 16 06:56:47 -0.528819 0.059676
BFGS: 17 06:56:50 -0.530000 0.116608
BFGS: 18 06:56:52 -0.530997 0.134195
BFGS: 19 06:56:55 -0.533622 0.083857
BFGS: 20 06:56:58 -0.535061 0.014591
Step Time Energy fmax
BFGS: 0 06:57:00 -0.407025 1.892326
BFGS: 1 06:57:03 -0.408503 2.303505
BFGS: 2 06:57:06 -0.441894 0.786452
BFGS: 3 06:57:08 -0.466190 0.758769
BFGS: 4 06:57:11 -0.513439 1.207292
BFGS: 5 06:57:13 -0.546304 0.396713
BFGS: 6 06:57:16 -0.547521 0.180180
BFGS: 7 06:57:19 -0.547612 0.316187
BFGS: 8 06:57:21 -0.548640 0.369188
BFGS: 9 06:57:24 -0.554280 0.311177
BFGS: 10 06:57:26 -0.557962 0.157095
BFGS: 11 06:57:29 -0.559328 0.087108
BFGS: 12 06:57:32 -0.560009 0.184298
BFGS: 13 06:57:34 -0.560848 0.233146
BFGS: 14 06:57:37 -0.561813 0.182305
BFGS: 15 06:57:39 -0.562527 0.179052
BFGS: 16 06:57:42 -0.564348 0.066446
BFGS: 17 06:57:45 -0.565037 0.053447
BFGS: 18 06:57:47 -0.565820 0.080173
BFGS: 19 06:57:50 -0.566627 0.111208
BFGS: 20 06:57:52 -0.567456 0.085274
BFGS: 21 06:57:55 -0.567385 0.041022
Step Time Energy fmax
BFGS: 0 06:57:58 -0.183217 0.913307
BFGS: 1 06:58:00 -0.213251 1.217160
BFGS: 2 06:58:03 -0.382715 1.929758
BFGS: 3 06:58:05 -0.156576 4.867807
BFGS: 4 06:58:08 -0.438516 0.965819
BFGS: 5 06:58:11 -0.458081 0.561021
BFGS: 6 06:58:13 -0.474265 0.719638
BFGS: 7 06:58:16 -0.478837 0.440359
BFGS: 8 06:58:19 -0.488117 0.285292
BFGS: 9 06:58:21 -0.500716 0.196141
BFGS: 10 06:58:24 -0.505062 0.264221
BFGS: 11 06:58:27 -0.506723 0.225622
BFGS: 12 06:58:29 -0.508863 0.062135
BFGS: 13 06:58:32 -0.510002 0.119993
BFGS: 14 06:58:35 -0.512013 0.222078
BFGS: 15 06:58:37 -0.513899 0.242976
BFGS: 16 06:58:40 -0.516163 0.220272
BFGS: 17 06:58:42 -0.518025 0.147682
BFGS: 18 06:58:45 -0.518703 0.053366
BFGS: 19 06:58:48 -0.519164 0.090695
BFGS: 20 06:58:50 -0.519974 0.124676
BFGS: 21 06:58:53 -0.521653 0.123972
BFGS: 22 06:58:55 -0.522448 0.067766
BFGS: 23 06:58:58 -0.523590 0.030779
Step Time Energy fmax
BFGS: 0 06:59:01 -0.183727 0.796283
BFGS: 1 06:59:03 -0.201260 0.721003
BFGS: 2 06:59:06 -0.280543 3.804196
BFGS: 3 06:59:09 -0.346087 0.741638
BFGS: 4 06:59:11 -0.377771 0.655889
BFGS: 5 06:59:14 -0.399061 0.553776
BFGS: 6 06:59:17 -0.406788 0.459194
BFGS: 7 06:59:19 -0.439583 0.613222
BFGS: 8 06:59:22 -0.457373 0.455756
BFGS: 9 06:59:24 -0.460248 0.186798
BFGS: 10 06:59:27 -0.465336 0.232088
BFGS: 11 06:59:30 -0.475927 0.407241
BFGS: 12 06:59:32 -0.481505 0.162820
BFGS: 13 06:59:35 -0.482997 0.143595
BFGS: 14 06:59:38 -0.485285 0.153875
BFGS: 15 06:59:40 -0.486864 0.065141
BFGS: 16 06:59:43 -0.487743 0.086161
BFGS: 17 06:59:45 -0.487887 0.107434
BFGS: 18 06:59:48 -0.488755 0.099689
BFGS: 19 06:59:51 -0.488699 0.060947
BFGS: 20 06:59:53 -0.489332 0.028248
Step Time Energy fmax
BFGS: 0 06:59:56 -0.212047 1.031853
BFGS: 1 06:59:58 -0.243991 1.358890
BFGS: 2 07:00:01 -0.372784 1.992701
BFGS: 3 07:00:04 -0.359998 2.287164
BFGS: 4 07:00:06 -0.460131 0.882945
BFGS: 5 07:00:09 -0.480119 0.445153
BFGS: 6 07:00:12 -0.490559 0.293096
BFGS: 7 07:00:14 -0.492229 0.242251
BFGS: 8 07:00:17 -0.505791 0.170451
BFGS: 9 07:00:20 -0.510150 0.216245
BFGS: 10 07:00:22 -0.511888 0.186415
BFGS: 11 07:00:25 -0.513011 0.082089
BFGS: 12 07:00:28 -0.515313 0.213305
BFGS: 13 07:00:30 -0.517419 0.289805
BFGS: 14 07:00:33 -0.520245 0.297473
BFGS: 15 07:00:36 -0.521741 0.251456
BFGS: 16 07:00:38 -0.522436 0.139359
BFGS: 17 07:00:41 -0.523054 0.105067
BFGS: 18 07:00:44 -0.524129 0.156845
BFGS: 19 07:00:46 -0.526990 0.155614
BFGS: 20 07:00:49 -0.529207 0.074389
BFGS: 21 07:00:52 -0.530374 0.016279
Step Time Energy fmax
BFGS: 0 07:00:54 -0.209529 1.196351
BFGS: 1 07:00:57 -0.218996 1.409690
BFGS: 2 07:01:00 -0.253565 1.243833
BFGS: 3 07:01:02 -0.337129 3.237358
BFGS: 4 07:01:05 -0.400880 0.639518
BFGS: 5 07:01:07 -0.423946 0.355679
BFGS: 6 07:01:10 -0.438820 0.565421
BFGS: 7 07:01:13 -0.447446 0.513742
BFGS: 8 07:01:15 -0.477670 0.207321
BFGS: 9 07:01:18 -0.481214 0.286810
BFGS: 10 07:01:21 -0.486647 0.231519
BFGS: 11 07:01:23 -0.491536 0.291956
BFGS: 12 07:01:26 -0.497632 0.296611
BFGS: 13 07:01:29 -0.502874 0.213218
BFGS: 14 07:01:31 -0.506446 0.137932
BFGS: 15 07:01:34 -0.509690 0.147959
BFGS: 16 07:01:37 -0.511731 0.199082
BFGS: 17 07:01:39 -0.516014 0.352286
BFGS: 18 07:01:42 -0.519890 0.243780
BFGS: 19 07:01:45 -0.521899 0.130866
BFGS: 20 07:01:47 -0.524292 0.117220
BFGS: 21 07:01:50 -0.527171 0.175774
BFGS: 22 07:01:52 -0.528817 0.258221
BFGS: 23 07:01:55 -0.532661 0.236111
BFGS: 24 07:01:58 -0.535161 0.135796
BFGS: 25 07:02:00 -0.536619 0.115409
BFGS: 26 07:02:03 -0.537893 0.155914
BFGS: 27 07:02:06 -0.539680 0.208948
BFGS: 28 07:02:08 -0.544244 0.168045
BFGS: 29 07:02:11 -0.548011 0.095706
BFGS: 30 07:02:14 -0.550467 0.084526
BFGS: 31 07:02:16 -0.551242 0.112795
BFGS: 32 07:02:19 -0.552897 0.107355
BFGS: 33 07:02:22 -0.555253 0.110939
BFGS: 34 07:02:24 -0.557344 0.108789
BFGS: 35 07:02:27 -0.558961 0.094271
BFGS: 36 07:02:30 -0.559958 0.109794
BFGS: 37 07:02:32 -0.560215 0.081836
BFGS: 38 07:02:35 -0.559996 0.076626
BFGS: 39 07:02:37 -0.561557 0.099462
BFGS: 40 07:02:40 -0.561613 0.086527
BFGS: 41 07:02:43 -0.562804 0.048196
Parse the trajectories and post-process#
As a post-processing step we check to see if:
the adsorbate desorbed
the adsorbate disassociated
the adsorbate intercalated
the surface has changed
We check these because they effect our referencing scheme and may result in erroneous energies. For (4), the relaxed surface should really be supplied as well. It will be necessary when correcting the SP / RX energies later. Since we don’t have it here, we will ommit supplying it, and the detector will instead compare the initial and final slab from the adsorbate-slab relaxation trajectory. If a relaxed slab is provided, the detector will compare it and the slab after the adsorbate-slab relaxation. The latter is more correct! Note: for the results in the AdsorbML paper, we did not check if the adsorbate was intercalated (is_adsorbate_intercalated()) because it is a new addition.
# Iterate over trajs to extract results
results = []
for file in glob(f"data/{bulk}_{adsorbate}/*.traj"):
rx_id = file.split("/")[-1].split(".")[0]
traj = ase.io.read(file, ":")
# Check to see if the trajectory is anomolous
initial_atoms = traj[0]
final_atoms = traj[-1]
atom_tags = initial_atoms.get_tags()
detector = DetectTrajAnomaly(initial_atoms, final_atoms, atom_tags)
anom = (
detector.is_adsorbate_dissociated()
or detector.is_adsorbate_desorbed()
or detector.has_surface_changed()
or detector.is_adsorbate_intercalated()
)
rx_energy = traj[-1].get_potential_energy()
results.append({"relaxation_idx": rx_id, "relaxed_atoms": traj[-1],
"relaxed_energy_ml": rx_energy, "anomolous": anom})
df = pd.DataFrame(results)
df
| relaxation_idx | relaxed_atoms | relaxed_energy_ml | anomolous | |
|---|---|---|---|---|
| 0 | 1 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.523903 | False |
| 1 | 7 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.562804 | False |
| 2 | 0 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.565065 | False |
| 3 | 5 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.489332 | False |
| 4 | 2 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.535061 | False |
| 5 | 3 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.567385 | False |
| 6 | 4 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.523590 | False |
| 7 | 6 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.530374 | False |
#scrap anomalies
df = df[~df.anomolous].copy().reset_index()
(Optional) Deduplicate structures#
We may have enumerated very similar structures or structures may have relaxed to the same configuration. For this reason, it is advantageous to cull systems if they are very similar. This results in marginal improvements in the recall metrics we calculated for AdsorbML, so it wasnt implemented there. It is, however, a good way to prevent wasteful VASP calculations. You can also imagine that if we would have enumerated 1000 configs per slab adsorbate combo rather than 100 for AdsorbML, it is more likely that having redundant systems would reduce performance, so its a good thing to keep in mind. This may be done by eye for a small number of systems, but with many systems it is easier to use an automated approach. Here is an example of one such approach, which uses a SOAP descriptor to find similar systems.
# Extract the configs and their energies
def deduplicate(configs_for_deduplication: list,
adsorbate_binding_index: int,
cosine_similarity = 1e-3,
):
"""
A function that may be used to deduplicate similar structures.
Among duplicate entries, the one with the lowest energy will be kept.
Args:
configs_for_deduplication: a list of ML relaxed adsorbate-
surface configurations.
cosine_similarity: The cosine simularity value above which,
configurations are considered duplicate.
Returns:
(list): the indices of configs which should be kept as non-duplicate
"""
energies_for_deduplication = np.array([atoms.get_potential_energy() for atoms in configs_for_deduplication])
# Instantiate the soap descriptor
soap = SOAP(
species=np.unique(configs_for_deduplication[0].get_chemical_symbols()),
r_cut = 2.0,
n_max=6,
l_max=3,
periodic=True,
)
#Figure out which index cooresponds to
ads_len = list(configs_for_deduplication[0].get_tags()).count(2)
position_idx = -1*(ads_len-adsorbate_binding_index)
# Iterate over the systems to get the SOAP vectors
soap_desc = []
for config in configs_for_deduplication:
soap_ex = soap.create(config, centers=[position_idx])
soap_desc.extend(soap_ex)
soap_descs = np.vstack(soap_desc)
#Use euclidean distance to assess similarity
distance = squareform(pdist(soap_descs, metric="cosine"))
bool_matrix = np.where(distance <= cosine_similarity, 1, 0)
# For configs that are found to be similar, just keep the lowest energy one
idxs_to_keep = []
pass_idxs = []
for idx, row in enumerate(bool_matrix):
if idx in pass_idxs:
continue
elif sum(row) == 1:
idxs_to_keep.append(idx)
else:
same_idxs = [row_idx for row_idx, val in enumerate(row) if val == 1]
pass_idxs.extend(same_idxs)
# Pick the one with the lowest energy by ML
min_e = min(energies_for_deduplication[same_idxs])
idxs_to_keep.append(list(energies_for_deduplication).index(min_e))
return idxs_to_keep
configs_for_deduplication = df.relaxed_atoms.tolist()
idxs_to_keep = deduplicate(configs_for_deduplication, adsorbate.binding_indices[0])
# Flip through your configurations to check them out (and make sure deduplication looks good)
print(idxs_to_keep)
view_x3d_n(configs_for_deduplication[2].repeat((2,2,1)))
df = df.iloc[idxs_to_keep]
low_e_values = np.round(df.sort_values(by = "relaxed_energy_ml").relaxed_energy_ml.tolist()[0:5],3)
print(f"The lowest 5 energies are: {low_e_values}")
df
The lowest 5 energies are: [-0.567 -0.535 -0.524 -0.489]
| index | relaxation_idx | relaxed_atoms | relaxed_energy_ml | anomolous | |
|---|---|---|---|---|---|
| 0 | 0 | 1 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.523903 | False |
| 5 | 5 | 3 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.567385 | False |
| 3 | 3 | 5 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.489332 | False |
| 4 | 4 | 2 | (Atom('Cu', [-1.3000465215529715, 2.2517466275... | -0.535061 | False |
Write VASP input files#
This assumes you have access to VASP pseudopotentials and the right environment variables configured for ASE. The default VASP flags (which are equivalent to those used to make OC20) are located in ocdata.utils.vasp. Alternatively, you may pass your own vasp flags to the write_vasp_input_files function as vasp_flags.
# Grab the 5 systems with the lowest energy
configs_for_dft = df.sort_values(by = "relaxed_energy_ml").relaxed_atoms.tolist()[0:5]
config_idxs = df.sort_values(by = "relaxed_energy_ml").relaxation_idx.tolist()[0:5]
# Write the inputs
for idx, config in enumerate(configs_for_dft):
os.mkdir(f"data/{config_idxs[idx]}")
write_vasp_input_files(config, outdir = f"data/{config_idxs[idx]}/")