.pth file for Nequip-OAM-L model

[dominicvarghese]

Hi all,

Where can we get the .pth or .pt file for Nequip-OAM-L model to use in LAMMPS?

[cw-tan]

Download the .nequip.zip file from https://www.nequip.net/models/mir-group/NequIP-OAM-L:0.1. See https://nequip.readthedocs.io/en/latest/guide/getting-started/files.html#package-files for details on package files.

Then see https://nequip.readthedocs.io/en/latest/integrations/lammps/index.html for LAMMPS. You could use pair_nequip or ML-IAP. You have to compile/prepare the packaged model file for use in LAMMPS.

[dominicvarghese]

Thanks for the reply. For fine-tuning the pre-trained NequIP model on a small, specific DFT dataset, what is the recommended practice for modifying the config.yaml?

My current approach is to use the tutorial.yaml (https://github.com/mir-group/nequip/blob/main/configs/tutorial.yaml) as a base and substitute the model definition with the ModelFromPackage loader. Besides updating the data path and reducing the learning rate, are there other essential modifications or best practices I should consider for adapting the training configuration for fine-tuning versus training from scratch?

[cw-tan]

That all sounds about right. You'll want to also change the per-type shifts to suit your specific DFT dataset's, i.e. compute the isolated atom energies for your DFT settings and update the shifts. Here's the docs for that https://nequip.readthedocs.io/en/latest/guide/training-techniques/fine_tuning.html

[dominicvarghese]

If a NequIP potential is trained exclusively on DFT data from 2x2x2 supercells, can it be expected to yield accurate energy and force predictions when used in LAMMPS for simulations with different supercell sizes (e.g., 3x3x3 or larger)?

[dominicvarghese]

So I was manually fine tuning the hyper parameters for the foundational model like the learning rate {0.01, 0.001, 0.0005}, weights in the loss function - energy:{10,50,100}, forces: {10,25,55}, stress:{1,5} and cutoff-radius {6,8,9}.

From these, the best I was able to get with this 446 configuration dataset - updated_dataset.txt using this yaml file - config.txt is as follows:
───────────────────────────────────
Validate metric DataLoader 0
───────────────────────────────────
val0_epoch/forces_mae 0.7656552427329895
val0_epoch/forces_rmse 8.666997710667008
val0_epoch/per_atom_energy_mae 0.415670453254575
val0_epoch/per_atom_energy_rmse 1.3103897269184126
val0_epoch/stress_mae 0.22452373458822372
val0_epoch/stress_rmse 0.4026654846852401
val0_epoch/total_energy_mae 16.626818130182997
val0_epoch/total_energy_rmse 52.4155890767365
val0_epoch/weighted_sum 12.327982665963496
────────────────────────────────────
Test metric DataLoader 0
───────────────────────────────────
test0_epoch/forces_mae 0.5284774567611029
test0_epoch/forces_rmse 4.990268992734253
test0_epoch/per_atom_energy_mae 0.42138581178916823
test0_epoch/per_atom_energy_rmse 0.6546921724867711
test0_epoch/stress_mae 0.2239569226918856
test0_epoch/stress_rmse 0.391528192840677
test0_epoch/total_energy_mae 16.85543247156673
test0_epoch/total_energy_rmse 26.187686899470844
test0_epoch/weighted_sum 7.643299636029939

I used this model to get the lattice relaxation on KTaO3 2x2x2 supercell, but the MD run shrinked the supercell. Initializing it from 7.974 A for the supercell, it just was shrinking to 5.4 A as seen from this log file - lammps.txt.

Are there any suggestions to further improve this model other than to add more data for the training part?

Also is there a way to automate the hyperparamter tuning in the yaml file instead of manually making runs?

[cw-tan]

Note that the cutoff radius of the foundation potential is fixed. There's nothing you can do during fine-tuning to change that. Presumably, you're just changing the cutoff radius in the dataset transform that creates the neighborlist if you're changing that variable during fine-tuning. That transform should be using the cutoff radius corresponding to the foundation potential.

Your errors are still quite large. Hard to tell if it's a dataset problem or a training problem without seeing your training curves. For hyperparameter tuning consider using https://nequip.readthedocs.io/en/latest/guide/training-techniques/wandb_example.html. In general, for training, it is useful to log the metrics on something like weights and biases or tensorboard.

[kavanase]

Agreed with @cw-tan , your errors are very very large, as also evident from the energy volume curves you showed, so there is likely some fundamental issue with your data/training setup here.
I think figuring out the issue with from-scratch training first, might be the most efficient approach to figuring out your issue.

[dominicvarghese]

I think this seems to be something with the training data that I am using. I used the foundational model for BaO with the following fine-tuning training dataset:

1. relaxed structure and isotropic strain relaxed structures (~20)
2. NPT at 300K for each of the strain (~300)
3. 10 random perturbation for each of the isotropic strain states (~100)

The metrics obtained were as follows:
────────────────────────────────────
Validate metric DataLoader 0
────────────────────────────────────
val0_epoch/forces_mae 0.035467296472225684
val0_epoch/forces_rmse 0.049854728093739656
val0_epoch/per_atom_energy_mae 0.006296619327083802
val0_epoch/per_atom_energy_rmse 0.010977016585429546
val0_epoch/stress_mae 0.0035850671552490566
val0_epoch/stress_rmse 0.00899625787540371
val0_epoch/total_energy_mae 0.10074590923334083
val0_epoch/total_energy_rmse 0.17563226536687274
val0_epoch/weighted_sum 0.05654320759352397
────────────────────────────────────
Test metric DataLoader 0
─────────────────────────────────────
test0_epoch/forces_mae 0.033611955947609136
test0_epoch/forces_rmse 0.052897130038291065
test0_epoch/per_atom_energy_mae 0.00880671828652072
test0_epoch/per_atom_energy_rmse 0.022027774789249756
test0_epoch/stress_mae 0.003699230525949872
test0_epoch/stress_rmse 0.011224639185390286
test0_epoch/total_energy_mae 0.14090749258433152
test0_epoch/total_energy_rmse 0.3524443966279961
test0_epoch/weighted_sum 0.0925156038701329

This seems to be better than for the previous case of KTaO3.
For this model I also tried plotting the Energy vs Volume curve as follows:

So I think the issue with the KTaO3 case was the training dataset(I guess).

[dominicvarghese]

I updated the dataset used for KTaO3 and the model seems to fit well with the DFT data. The metrics are as follows:
───────────────────────────────────
Validate metric DataLoader 0
───────────────────────────────────
val0_epoch/forces_mae 0.03797666172679939
val0_epoch/forces_rmse 0.0541983196820138
val0_epoch/per_atom_energy_mae 0.0030298049687999926
val0_epoch/per_atom_energy_rmse 0.00490015769958846
val0_epoch/stress_mae 0.00245478106041206
val0_epoch/stress_rmse 0.005799939842214366
val0_epoch/total_energy_mae 0.12119219875200152
val0_epoch/total_energy_rmse 0.19600630798354093
val0_epoch/weighted_sum 0.06020440160915106
───────────────────────────────────────
Test metric DataLoader 0
───────────────────────────────────────
test0_epoch/forces_mae 0.04018132698504112
test0_epoch/forces_rmse 0.05819380972603791
test0_epoch/per_atom_energy_mae 0.0031531600522605174
test0_epoch/per_atom_energy_rmse 0.005844569740517788
test0_epoch/stress_mae 0.002565498609153855
test0_epoch/stress_rmse 0.005705786717596827
test0_epoch/total_energy_mae 0.1261264020904147
test0_epoch/total_energy_rmse 0.23378278962070945
test0_epoch/weighted_sum 0.06734538583976528
────────────────────────────────────────────
Plotting energy vs volume plot gave me this:

So I was using this model to look at the lattice parameter through ASE by using this python script - lattice.py.txt and gave me the unitcell lattice parameter of the cubic state to be 3.9938 A (pretty close to the experimental value).

But when I try to use the LAMMPS version using this input script - in.lammps_relax.txt, I get the lattice to just increase in size as seen from the log file - log_lammps.txt. The 2x2x2 supercell which is supposed to be ~ 8 A in length goes to 24 A and keeps on increasing in volume.

Is it because I am making some mistake in the LAMMPS input script I am using?

[dominicvarghese]

Do we include all the different configurations involved in the ionic relaxation process to the training set (does all the configurations from the vasprun.xml file is needed to make the .xyz file required for training)? Or do we just need the final configuration and its energy, stress values?

[cw-tan]

if it's well-converged, good quality DFT, it's useful information for training the model. Though some sampling scheme might be useful (or not), depends.

[dominicvarghese]

I was doing hyperparameter tuning with the sweep file sweep.yaml for the following config file and dataset.

I was varying the following parameters:

  learning_rate:
    values: [5e-4,1e-3]

  loss_coeff_energy:
    values: [1, 10, 20]

  loss_coeff_forces:
    values: [1, 5]

This is how I was getting the validation epoch metrics from wandb

From the figure, energy and force seems to be having a difficult time with the validation part compared to stress. Are there any suggestion to make the model better?

The following is the free-energy distribution of the dataset used:

[dominicvarghese]

So with this code,

calculator = NequIPCalculator.from_compiled_model(
    compile_path=compile_path,
    device="cuda" if torch.cuda.is_available() else "cpu",
    chemical_symbols=["Ba", "O"],
)

chemical symbols tag needs to be removed? Just to add, I used to get nearly accurate lattice constants when I was using just elements under consideration in the chemical_symbols tag in the config.yaml and then using the output with the ASE calculator.

The LAMMPS version should work fine then?

[cw-tan]

chemical symbols tag needs to be removed?
I believe so, it'll throw a warning but you can ignore it.

Just to add, I used to get nearly accurate lattice constants when I was using just elements under consideration in the chemical_symbols tag in the config.yaml and then using the output with the ASE calculator.

This is for fine-tuning? That sounds right if the training process managed to train the pre-trained embeddings for other atomic environment, like ["H", "He"] if you previously used ["Ba", "O"], to your present system. But it won't be fully leveraging on the Ba and O related data the original foundation potential was trained on. i.e. there's a mismatch between the atomic species, unless you specify the full set of chemical symbols in the exact same order as the ones used during training of the foundation potential. This same problem happens with the ASE calculator when you specify

calculator = NequIPCalculator.from_compiled_model(
    compile_path=compile_path,
    device="cuda" if torch.cuda.is_available() else "cpu",
    chemical_symbols=["Ba", "O"],
)

if you're using the foundation potential directly. The model will be using the embeddings for "H" and "He" (the first two entries) to treat "Ba" and "O" atoms.

[cw-tan]

LAMMPS should be fine since you are mapping to atom type names directly

[dominicvarghese]

I used the full chemical symbols tag in the config.yaml file and removed the chemical symbols argument from the ASE calculator part. The metrics are as follows:

          Validate metric                 DataLoader 0
    val0_epoch/forces_mae          0.015034316562774395
   val0_epoch/forces_rmse           0.02919918694481102
val0_epoch/per_atom_energy_mae      0.0035691677766850997
val0_epoch/per_atom_energy_rmse     0.005437235846266573
    val0_epoch/stress_mae          0.0019653804297617753
   val0_epoch/stress_rmse          0.0032203667823550385
 val0_epoch/total_energy_mae       0.057106684426961596
val0_epoch/total_energy_rmse        0.08699577354026516
   val0_epoch/weighted_sum         0.025316014038735082

         Test metric                     DataLoader 0
    test0_epoch/forces_mae           0.015465030313636614
   test0_epoch/forces_rmse           0.03163172230887766
test0_epoch/per_atom_energy_mae      0.0035840797003043267
test0_epoch/per_atom_energy_rmse      0.005468176775181235
    test0_epoch/stress_mae          0.0019349161843746003
   test0_epoch/stress_rmse          0.0031234462643849143
 test0_epoch/total_energy_mae        0.05734527520486923
test0_epoch/total_energy_rmse        0.08749082840289976
   test0_epoch/weighted_sum          0.025755434394316046

I need to do some hyperparameter tuning on this (errors aren't that bad for now).
The energy vs volume curve looks like

I tried with both ASE and LAMMPS to get the relaxed lattice parameter.
With ASE, the lattice parameter came to be 5.53713 A and with LAMMPS, the relaxed unitcell had lattice parameter of 5.53679 A.
The DFT relaxed unitcell had lattice parameter of 5.53257 A (difference less than 0.1%). So I think the fine-tuning on the pre-trained model went well. Thanks for the support throughout this.

[cw-tan]

Nice, thank you for being thorough and helping us find pitfalls in our fine-tuning infrastructure. I'll add more user documentation and make minor breaking API changes for future safety and to reduce confusion.