Implement ZBL potential

[peastman]

This is the first piece of #26. It isn't fully tested yet, but I think it's ready for a first round of comments.

I created a mechanism for passing arguments to the prior. prior_args is now the option specified by the user in the config file. prior_init_args stores the value returned by get_init_args(), which contains the arguments needed for reconstructing it from a checkpoint.

This prior requires the dataset to provide several pieces of information. To keep things general, the HDF5 format allows the file to contain a _metadata group which can store arbitrary pieces of information. Most of the other dataset classes should be able to hardcode the necessary values, since they aren't intended to be general.

[stefdoerr]

Looks good to me. Tests are failing though:

 E       RuntimeError: 
E       'NoneType' object has no attribute or method 'atomwise'.:
E         File "/usr/share/miniconda/envs/torchmd-net/lib/python3.9/site-packages/torchmdnet/models/model.py", line 182
E           
E               # apply atomwise prior model
E               if self.prior_model is not None and self.prior_model.atomwise:
E                                                   ~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
E                   x = self.prior_model(x, z, pos, batch)

Also you might have some merge conflicts with the latest PR of Raimondas

[stefdoerr]

Wait a second, how does this error even occur? Shouldn't the if short-circuit if self.prior_model is not None? Weird. Maybe jit-ed conditionals don't short-circuit?

[giadefa]

Who runs the model knows what units want to use and set it for the priors

On Fri, Oct 7, 2022, 16:28 Peter Eastman @.***> wrote:

@.**** commented on this pull request.

In torchmdnet/priors.py https://github.com/torchmd/torchmd-net/pull/134#discussion_r990477979:

@@ -76,3 +79,52 @@ def get_init_args(self):

 def forward(self, x, z, pos, batch):
     return x + self.atomref(z)

+ + +class ZBL(BasePrior):

• """This class implements the Ziegler-Biersack-Littmark (ZBL) potential for screened nuclear repulsion.
• Is is described in https://doi.org/10.1007/978-3-642-68779-2_5 (equations 9 and 10 on page 147). It
• is an empirical potential that does a good job of describing the repulsion between atoms at very short
• distances.
• To use this prior, the Dataset must provide the following attributes.
• atomic_number: 1D tensor of length max_z. atomic_number[z] is the atomic number of atoms with atom type z.
• distance_scale: multiply by this factor to convert coordinates stored in the dataset to meters

The neural network is agnostic with regard to units. It receives inputs and tries to make its outputs match the ones in the dataset. It doesn't know or care what units they're in.

A physics based potential does depend on units. It needs to know what units the coordinates are in, and it needs to know what units the energy is expected to be in.

— Reply to this email directly, view it on GitHub https://github.com/torchmd/torchmd-net/pull/134#discussion_r990477979, or unsubscribe https://github.com/notifications/unsubscribe-auth/AB3KUOUMNVJC6NDOJTC7W7LWCCBXVANCNFSM6AAAAAAQ67RI4M . You are receiving this because you commented.Message ID: @.***>

[peastman]

Who runs the model knows what units want to use and set it for the priors

The units are determined by the dataset. It isn't a free choice for the user to make. If your dataset contains positions in nm and energies in kJ/mol, that's what the prior needs to work with. Any other units would produce wrong results.

We could consider creating an automated unit conversion system like I suggested in https://github.com/torchmd/torchmd-net/issues/26#issuecomment-1266154668, but that's a separate project.

[peastman]

Wait a second, how does this error even occur? Shouldn't the if short-circuit if self.prior_model is not None?

Yes it should! I think it's a bug in the torchscript jit. I was able to work around it by splitting it into two lines.

[peastman]

This is ready for review.

[stefdoerr]

The test is failing if you could take a look

[peastman]

The test is failing if you could take a look

Fixed. I couldn't reproduce it locally, but I managed to guess what it was unhappy about.

[peastman]

I added the multiple priors support. I tried to make the syntax fairly flexible. All of the following are valid.

prior_model: Atomref

prior_model:
  - ZBL:
      cutoff_distance: 4.0
      max_num_neighbors: 50
  - Atomref

prior_model:
  - ZBL:
      cutoff_distance: 4.0
      max_num_neighbors: 50
    Atomref:

I added tests for creating models both from config files and from checkpoints. And the test for older checkpoints in test_module.py continues to work. I can't be sure I've caught all possible problems though, so if you can check your own files that would be helpful.

[peastman]

Is this ok to merge?