mir-group / nequip

NequIP is a code for building E(3)-equivariant interatomic potentials
https://www.nature.com/articles/s41467-022-29939-5
MIT License
611 stars 135 forks source link

Allow calculation of isolated atom #426

Closed tjgiese closed 4 months ago

tjgiese commented 5 months ago

Description

  1. plot_dimers.py: The --n-samples was not being respected when calculating the linear spacing. Furthermore, if the model contained a single atom type, then matplotlib returns a scalar axes variable rather than a list.
  2. _atomwise.py: Prevent assert failure when the system consists of a single atom
  3. _grad_output.py: If a system consists of a single atom, then pytorch will complain that there are no gradients to calculate. The proposed "fix" is to check if all of the configurations consist of 1 atom; if so, then don't call autograd - instead return zeros.

Motivation and Context

Previous version of nequip could not calculate the energy of an isolate atom. One often wants to create loss functions that fit relative energies (not just total energies). If an isolated atom cannot be calculated, then it prevents supermolecular calculation of molecule-atomic ion interactions, for example. A discussion of this issue, and the calculation of well-separated dimers, can be found on the allegro github project: https://github.com/mir-group/allegro/discussions/82

How Has This Been Tested?

I followed the nequip/allegro tutorial to generate si-deployed.pth. I then created a python script that reads the model file and an ase xyz file. The model is used to calculate the energy and forces. The energy of a well-separated dimer is verified to be twice the energy of the monomer.

Test environment

Types of changes

Checklist: