Constrained Atoms

[tgmaxson]

I am curious about how constrained atoms are handled in NEQUIP. For example, I have atoms which are fully constrained in XYZ that should not be included in the model. Are these currently trained into the model and/or is there a way to filter them out?

[Linux-cpp-lisp]

Hi @tgmaxson ,

You'll have to clarify what you mean by "not included in the model"--- do you want to omit them from the input? Made predictions on them but not train on them? Etc.

[tgmaxson]

I would like to not train on them and there is no real reason to make predictions on them either (they will be 0,0,0 in force). Omitting them from the total energy/stress should also be fine since they will remain constant. However, they need to be there for the other atoms.

To give a more complete context, this is something like a FCC111 surface with 2 layers constrained at the bottom and 2 layers relaxed.

What are the relevant tags for this or am I overthinking this and they should just stay in the model?

[Linux-cpp-lisp]

In general I have understood constrained atoms to be zeros in velocity (as in, you don't time integrate them), but not zero in force. (Though obviously you could enforce zero velocity by setting forces to zero.) I definitely agree that if in your data the forces are fake zeros then you should not train on that, since it will screw everything up 😁 If those atoms have real forces in your original data, it doesn't matter whether they are constrained or not and you can just use that as more data-- don't have to do anything different.

NequIP does suppor this: It's not something I've worked with recently but basically if you provide NaNs as the target/truth value for specific atoms they will be correctly ignored in loss calculations.

[tgmaxson]

To set the target/truth value I simply need to set the values to NaN then right? It would be great if there was a way in NEQUIP to do this automagically. To be clear, just set the force values in the supplied dataset to NaN.

Constraints do typically mean zero velocity for MD / optimization but typically codes such as VASP will report it as (0, 0, 0) in force as well. It's not ideal.

[Linux-cpp-lisp]

Right, nequip does do this automagically if you just provide the targets for those atoms as NaN. You can confirm this by setting all the targets to NaN and seeing that the loss will be zero rather than NaN.

...VASP will report it as (0, 0, 0) in force as well. It's not ideal.

Ah I see, yes that is annoying but I see how that limits what you can do here.

[tgmaxson]

Is there anyway you can confirm this works? I have slightly gone down the rabbit hole looking for why I am having an issue in my model and discovered two things seperately.

ASE's "get_forces()" function will convert np.nan to 0 (which is an ASE bug) which means nothing will work

Looking at nequip.train._loss() it looks like "ignore_nan" defaults to False, which makes me think this won't work by default. Can you somehow confirm that setting forces to np.nan will infact not include it in the loss function?

[Linux-cpp-lisp]

@nw13slx how did you do this in the past?

I'm honestly not sure if this works at the moment with an ASE dataset (which is unhelpful, I know, since usually we encourage using that). If all your frames have the same number of atoms, you can use the NPZ dataset, which will definitely handle NaN correctly.

The second issue though is unclear to me how to manually check.

Besides "print debugging", you can try making everything NaN and see if it all comes out as zeros.

Looking at nequip.train._loss() it looks like "ignore_nan" defaults to False, which makes me think this won't work by default.

Hm good catch, yes set ignore_nan: True in the YAML. You can add a print to the class you mentioned to confirm that this propagates through.

[tgmaxson]

For the moment I am simply avoiding this issue, however would it be trivial to patch the code to read constraints as NaN for loss function using ASE? This function can be used to check for constrained atoms (any constraint on an atom will give unphysical forces in my opinion and should be automatically removed, but you can decide on a default).

https://gitlab.com/ase/ase/-/blob/master/ase/constraints.py#L34

[Linux-cpp-lisp]

Interesting point, I was unaware of this "constrained atoms" data in ASE. Is there a way to specify it in extXYZ files? Or would we have to write that up ourselves...

[tgmaxson]

So typically ASE readers will import constraints with whatever the dataset is. So in the case of POSCARs, some atoms will come in with FixedAtom constraints for example.

Is the need for extXYZ internal to NEQUIP? its not elegant, but since I don't think extXYZ supports constraints correctly in ASE right now you could probably dump a really basic pickle/json file with just the constraints for each image. This can then either be put back together as constraints when loaded or kept as a simple array for zeroing out loss functions.

[Linux-cpp-lisp]

Is the need for extXYZ internal to NEQUIP?

Not at all--- NequIP deals directly with ase.Atoms, so in theory any format supported by ase is fine. In practice, we usually recommend and use extxyz because it is the best supported format for ase as well has having good extensibility to custom fields and data.

(In fact, with the exception of the initial preprocessing, NequIP also doesn't deal with ase.Atoms; it deals with it's own AtomicData class.)

I think the best way to deal with this, since it is a bit case specific, is with a custom AtomicInMemoryDataset subclass. See NPZDataset for a simple example of such a class; you'd want to make a copy of that basically that changes get_data() to use ase.io.read() to get your ase.Atoms() objects, AtomicData.from_ase to turn them into AtomicData, and then write in the relevant NaNs based on whatever format for the constraint data you choose to use. Then you just return a list of AtomicData from get_data().

To use this custom dataset, in your YAML, instead of

dataset: ase

you use

dataset: !!python/name:name.of.my.module.MyCustomDataset ''

[Linux-cpp-lisp]

For future reference, there now exists a BETA interface for using this that includes an example here:

https://github.com/mir-group/nequip/tree/masks/examples/mask_labels

on the masks branch off develop. If you find any issues with it, please feel free to let me know here!

[Linux-cpp-lisp]

See also tracking issue #307