Input files, periodic boundary conditions

[schiotz]

Hi,

I trying to familiarizing myself with nequip by repeating some simulations that a student did with other NN based potentials. In that connection, I looked at how to present the data for Nequip, and the NPZ files look like the most compact format. I tried to work out what to put in these files based on the examples, but have a few questions.

1. This data set has periodic boundary conditions. How am I supposed to specify that in the .npz file? I'll also need to specify the cell that is used. I can see that most fields are specified in the .yaml file, so I suppose I just need to add some field there.
2. How do I handle data sets with varying number of atoms? Typically it will be ab-initio MD runs with slightly different stochiometry. Can I specify multiple .npz files, or should I use an ASE .traj / ASE .db file?
3. How are training, validation and test sets handled? If I make a file named xyz-train.npz and xyz-test.npz, and specify the former in the .yaml file, will Nequip then split the -train.npz file into a training and a validation set when nequip-train is run? And will nequip-evaluate somehow extract the training file name from the .yaml file, and replace -train with -test? Or is something else happening behind the scenes?

Thank you very much for your help. I am hoping to soon start a student project using Nequip, so I better understand how it works first :-)

/Jakob

[Linux-cpp-lisp]

Hi Jakob,

Generally we recommend extxyz (ASE's favorite format) as the easiest and most flexible way to deal with data. In particular, it can handle your second point. The npz format only handles rectangular inputs where the number of atoms does not change, and we do not at this point have support for combining multiple datasets specified separately--- this is supported through combining them in preprocessing instead.

Point by point:

1. For npz, you provide a cell array of shape [n_frame, 3, 3] or [3, 3]. If the latter, you must also set

npz_fixed_field_keys:
  - cell

so that it knows that the same single cell is shared across all frames.
For extxyz, writing out the file through ASE will give correct results.

2. npz does not support this; extxyz supports it by default. As for ASE .traj/.db files, if they are supported transparently by ase.io.read, they can also be used directly by our ASE dataset loader. You may need to specify:

ase_args:
      format: extxyz

with extxyz replaced by the correct format string.

3. At no point do we "magically" determine sets based on file names. By default, the train and validation sets are specified by lists of frame indexes into the single dataset specified. By prepending the same YAML options with validation_, you also can specify an entirely different dataset from which to draw the validation frames.

Test sets are specified to nequip-evaluate when you call it, since they do not need to be determined in advance for training. By default, if you drew both train and validation frames from a single dataset for training, nequip-evaluate will use all other frames in that dataset as the test set. For more control, use the --dataset-config option to specify a YAML file that, using the same options as before, defines a dataset to use as the test set. (In that YAML file you can also specify specific frames to use from it as the test set.)

Let me know if this is unclear or incomplete, and looking forward to hearing about what you do with NequIP!

[schiotz]

Thanks, @Linux-cpp-lisp, I think I got most of it.

Regarding point 3: In the MRS example, we used the configs/example.yaml containing

n_train: 100                                                                       # number of training data
n_val: 50                                                                          # number of validation data
 [ ... ]
train_val_split: random

and a .npz file containing 1000 configurations (toluene_ccsd_t-train.npz). So in that case, 100 configs are picked at random for training, 50 for validation, and the remaining 850 for testing? And which ones were chosen are stored in some file that nequip-evaluate reads?

The test set seems a bit large, but I guess this is just an example. If I understood this right, then I think I know how to proceed.

Right now, I am trying to see how good a potential I can make using the data that we used in this publication: https://doi.org/10.1063/5.0077580

[Linux-cpp-lisp]

Right, exactly!

Re the size of the test set, it's not really that the remaining 850 are "chosen" as a test set. Only the train and validation indexes are stored. Rather, the use of the remaining frames is just an attempt to provide a semi-sane default test set that is disjoint with the train and validation sets. (In general, it seemed reasonable to use all available data to get the best estimate of your models accuracy; this is not appropriate of course if you are doing something fancier with your errors.)

If you want to choose a more limited test set from the same overall dataset, you can use the --test-indexes option to nequip-evaluate, which specifies a file that lists the indexes of all frames you wish to include in the test set. (Note that you are not prevented from including frames that were also used in the train/validation sets; if you want to avoid this, as you typically do, you can use train_idcs and val_idcs in config.yaml in your training directory to avoid choosing those frames.)

Full details can be found with

nequip-evaluate --help

Right now, I am trying to see how good a potential I can make using the data that we used in this publication: https://doi.org/10.1063/5.0077580

Very nice!

[schiotz]

Thank you very much for your informative answers! It is bed-time here in Denmark, so I'll let the GPUs hum for the night, and see tomorrow if it works :-)

[schiotz]

@Linux-cpp-lisp

Leaving the GPUs humming gave very promising results. Thanks!

I would like to be able to test different models, trained with different data sets, on the same test set. It would be nice if I could somehow specify a different file for nequip-evaluate, but I cannot see an easy way to do that. I guess I could copy the original .yaml file, replace the file name, set n_train and n_val to 0, and use --dataset-config to point to that file. Or would I need to edit the config.yaml file in the training folder?

Speaking of which: If I wanted to specify --test-indices, and wanted to read train_idcs and val_idcs from the config.yaml file, how would I do that? It seems to be some kind of opaque binary block of data. I tried reading it with the yaml module, but it gave the error
yaml.constructor.ConstructorError: could not determine a constructor for the tag 'tag:yaml.org,2002:python/object/apply:torch._utils._rebuild_tensor_v2'
As you may be able to see, I am new to both pytorch and yaml. :-)

Is there a Python interface to the operations of nequip? Sometimes that is the easiest way to do something fancy.

[Linux-cpp-lisp]

Great!

I would like to be able to test different models, trained with different data sets, on the same test set. It would be nice if I could somehow specify a different file for nequip-evaluate, but I cannot see an easy way to do that. I guess I could copy the original .yaml file, replace the file name, set n_train and n_val to 0, and use --dataset-config to point to that file. Or would I need to edit the config.yaml file in the training folder?

The copy is unnecessary--- if you specify a description of the test dataset in --dataset-config my-test-set.yaml, then it either uses the whole specified dataset (without --test-indexes) or the frames specified within that dataset by --test-indexes. No need for the rest of the original config, which gets picked up still from --train-dir.

yaml.constructor.ConstructorError: could not determine a constructor for the tag 'tag:yaml.org,2002:python/object/apply:torch._utils._rebuild_tensor_v2'

It's a serialized PyTorch tensor, so you just need to have import torched before you try to load the YAML.

Is there a Python interface to the operations of nequip? Sometimes that is the easiest way to do something fancy.

There is, especially for the lower-level operations. But for the "user-facing" commands, like nequip-evaluate, the "Python" interface is just the command line tool with the same arguments provided as a dict. If you really need to do something fancy, you can always look at the nequip-evaluate source code (nequip/scripts/evaluate.py), which uses APIs from other parts of the code, but I strongly recommend against doing that unless you truly need to. (Those internal APIs are changing unexpectedly, and there are a lot of ways to shoot yourself in the foot without knowing the internals of the internals 😄 )