Add alchemical group for FEP.

configs/nequix-omol-neutral-index.yml

@@ -0,0 +1,133 @@
+resume_from: nequix-omol-neutral-state.pkl
+state_path: nequix-omol-neutral-state.pkl # for saving/resuming training
+cutoff: 6.0
+
+# dataset specific
+train_path: "/cv/scratch/u/daigavaa/omol25/neutral_train"
+valid_path: "/cv/scratch/u/daigavaa/omol25/neutral_val"
+
+atomic_numbers: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83]
+
+# isolated atom energies (eV)
+# referenced from https://github.com/ACEsuit/mace-foundations/blob/main/mace_omol/omol-statistics-linear.json
+atom_energies:
+    1: -1.59900036e+01
+    2: -8.49660292e+01
+    3: -2.03766680e+02
+    4: -3.91507730e+02
+    5: -6.78593442e+02
+    6: -1.03716785e+03
+    7: -1.49037815e+03
+    8: -2.04759259e+03
+    9: -2.71737651e+03
+    10: -3511.7048
+    11: -4.41233388e+03
+    12: -5.44339111e+03
+    13: -6.58743663e+03
+    14: -7.88100913e+03
+    15: -9.29051601e+03
+    16: -1.08350869e+04
+    17: -1.25222590e+04
+    18: -1.43583656e+04
+    19: -1.63234949e+04
+    20: -1.84356504e+04
+    21: -2.06997835e+04
+    22: -2.31123879e+04
+    23: -2.56833420e+04
+    24: -2.84179222e+04
+    25: -3.13154372e+04
+    26: -3.43829680e+04
+    27: -3.76226929e+04
+    28: -4.10396553e+04
+    29: -4.46353620e+04
+    30: -4.84161616e+04
+    31: -5.23713649e+04
+    32: -5.65176633e+04
+    33: -6.08400888e+04
+    34: -6.53480263e+04
+    35: -7.00422412e+04
+    36: -7.49358992e+04
+    37: -6.53792089e+02
+    38: -8.32218179e+02
+    39: -1.04112032e+03
+    40: -1.27159339e+03
+    41: -1.54737085e+03
+    42: -1.84371065e+03
+    43: -2.18542972e+03
+    44: -2.57346296e+03
+    45: -3.00033628e+03
+    46: -3.47247059e+03
+    47: -3.99336867e+03
+    48: -4.56246488e+03
+    49: -5.16328012e+03
+    50: -5.82455062e+03
+    51: -6.53872001e+03
+    52: -7.29551322e+03
+    53: -8.09959609e+03
+    54: -8.96620607e+03
+    55: -5.46024853e+02
+    56: -6.90472676e+02
+    57: -8.55056742e+02
+    58: -1.29225322e+04
+    59: -1.40666521e+04
+    60: -1.52742278e+04
+    61: -1.65518215e+04
+    62: -1.79001435e+04
+    63: -1.93239888e+04
+    64: -2.08293417e+04
+    65: -2.24294711e+04
+    66: -2.40787048e+04
+    67: -2.57954236e+04
+    68: -2.76118340e+04
+    69: -2.95211005e+04
+    70: -3.15241178e+04
+    71: -3.36137406e+04
+    72: -1.29451208e+03
+    73: -1.54106188e+03
+    74: -1.81327813e+03
+    75: -2.11628877e+03
+    76: -2.45279159e+03
+    77: -2.82809875e+03
+    78: -3.23812286e+03
+    79: -3.68434882e+03
+    80: -4.16968620e+03
+    81: -4.68420895e+03
+    82: -5.24032637e+03
+    83: -5.83589325e+03
+
+avg_n_edges: 644.6255168799191
+avg_n_neighbors: 18.584025140193773
+avg_n_nodes: 27.072677525062158
+max_n_edges: 10046
+max_n_nodes: 181
+scale: 2.7086112502046458
+shift: 0.16233131071765744
+
+# model specific
+hidden_irreps: "128x0e + 64x1o + 32x2e + 32x3o"
+lmax: 3
+n_layers: 4
+radial_basis_size: 8
+radial_mlp_size: 64
+radial_mlp_layers: 2
+radial_polynomial_p: 6.0
+mlp_init_scale: 4.0
+index_weights: true
+layer_norm: true
+kernel: true
+
+# training specific
+optimizer: "muon"
+learning_rate: 0.01
+warmup_epochs: 0.1
+warmup_factor: 0.2
+grad_clip_norm: 100.0
+weight_decay: 1.0e-3
+batch_size: 128  # NOTE: per device
+n_epochs: 12
+energy_weight: 20.0
+force_weight: 20.0
+stress_weight: 0.0
+loss_type: "mae"
+log_every: 100
+ema_decay: 0.999

configs/nequix-omol-neutral.yml

@@ -0,0 +1,133 @@
+resume_from: nequix-omol-neutral-state.pkl
+state_path: nequix-omol-neutral-state.pkl # for saving/resuming training
+cutoff: 6.0
+
+# dataset specific
+train_path: "/cv/scratch/u/daigavaa/omol25/neutral_train"
+valid_path: "/cv/scratch/u/daigavaa/omol25/neutral_val"
+
+atomic_numbers: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83]
+
+# isolated atom energies (eV)
+# referenced from https://github.com/ACEsuit/mace-foundations/blob/main/mace_omol/omol-statistics-linear.json
+atom_energies:
+    1: -1.59900036e+01
+    2: -8.49660292e+01
+    3: -2.03766680e+02
+    4: -3.91507730e+02
+    5: -6.78593442e+02
+    6: -1.03716785e+03
+    7: -1.49037815e+03
+    8: -2.04759259e+03
+    9: -2.71737651e+03
+    10: -3511.7048
+    11: -4.41233388e+03
+    12: -5.44339111e+03
+    13: -6.58743663e+03
+    14: -7.88100913e+03
+    15: -9.29051601e+03
+    16: -1.08350869e+04
+    17: -1.25222590e+04
+    18: -1.43583656e+04
+    19: -1.63234949e+04
+    20: -1.84356504e+04
+    21: -2.06997835e+04
+    22: -2.31123879e+04
+    23: -2.56833420e+04
+    24: -2.84179222e+04
+    25: -3.13154372e+04
+    26: -3.43829680e+04
+    27: -3.76226929e+04
+    28: -4.10396553e+04
+    29: -4.46353620e+04
+    30: -4.84161616e+04
+    31: -5.23713649e+04
+    32: -5.65176633e+04
+    33: -6.08400888e+04
+    34: -6.53480263e+04
+    35: -7.00422412e+04
+    36: -7.49358992e+04
+    37: -6.53792089e+02
+    38: -8.32218179e+02
+    39: -1.04112032e+03
+    40: -1.27159339e+03
+    41: -1.54737085e+03
+    42: -1.84371065e+03
+    43: -2.18542972e+03
+    44: -2.57346296e+03
+    45: -3.00033628e+03
+    46: -3.47247059e+03
+    47: -3.99336867e+03
+    48: -4.56246488e+03
+    49: -5.16328012e+03
+    50: -5.82455062e+03
+    51: -6.53872001e+03
+    52: -7.29551322e+03
+    53: -8.09959609e+03
+    54: -8.96620607e+03
+    55: -5.46024853e+02
+    56: -6.90472676e+02
+    57: -8.55056742e+02
+    58: -1.29225322e+04
+    59: -1.40666521e+04
+    60: -1.52742278e+04
+    61: -1.65518215e+04
+    62: -1.79001435e+04
+    63: -1.93239888e+04
+    64: -2.08293417e+04
+    65: -2.24294711e+04
+    66: -2.40787048e+04
+    67: -2.57954236e+04
+    68: -2.76118340e+04
+    69: -2.95211005e+04
+    70: -3.15241178e+04
+    71: -3.36137406e+04
+    72: -1.29451208e+03
+    73: -1.54106188e+03
+    74: -1.81327813e+03
+    75: -2.11628877e+03
+    76: -2.45279159e+03
+    77: -2.82809875e+03
+    78: -3.23812286e+03
+    79: -3.68434882e+03
+    80: -4.16968620e+03
+    81: -4.68420895e+03
+    82: -5.24032637e+03
+    83: -5.83589325e+03
+
+avg_n_edges: 644.6255168799191
+avg_n_neighbors: 18.584025140193773
+avg_n_nodes: 27.072677525062158
+max_n_edges: 10046
+max_n_nodes: 181
+scale: 2.7086112502046458
+shift: 0.16233131071765744
+
+# model specific
+hidden_irreps: "128x0e + 64x1o + 32x2e + 32x3o"
+lmax: 3
+n_layers: 4
+radial_basis_size: 8
+radial_mlp_size: 64
+radial_mlp_layers: 2
+radial_polynomial_p: 6.0
+mlp_init_scale: 4.0
+index_weights: false
+layer_norm: true
+kernel: true
+
+# training specific
+optimizer: "muon"
+learning_rate: 0.01
+warmup_epochs: 0.1
+warmup_factor: 0.2
+grad_clip_norm: 100.0
+weight_decay: 1.0e-3
+batch_size: 128  # NOTE: per device
+n_epochs: 12
+energy_weight: 20.0
+force_weight: 20.0
+stress_weight: 0.0
+loss_type: "mae"
+log_every: 100
+ema_decay: 0.999