mir-group/nequip

๐ŸŒŸ [FEATURE] Custom properties

Opened this issue ยท 1 comments

Is your feature request related to a problem? Please describe.
Custom properties are supported. However, currently I encounter two issues:

  • The calculator class does not expose the custom properties.
  • For second order tensor predictions, if you use the current method for custom properties, you end up with an Irrep. for scalars and first orders, the irreps and cartesian tensors are the same (except maybe for some scalar multiplication which is irrelevant when fitting), but for higher order tensors the shape doesn't match.

Describe the solution you'd like

  • Some option added to the calculator class to acces the out dict of the model (e.g. self.out = out or a helper function out_transformer = lambda out, results: ... similar to how the input can be transformed.
  • Maybe a model builder can be added that transforms a tensor from irreps to explicit cartesian?

Describe alternatives you've considered

  • Reimplement the calculator class, but it requires copy pasting the entire calculate method since the output of the model is not accessible outside of it.
  • Implement the builder myself (which I've done, but maybe it's useful to make it available/make it more generic)

Hi @ThePauliPrinciple ,

Regarding "The calculator class does not expose the custom properties.", unless you are specifically running e.g. ASE MD, the easiest thing to do is just to skip the ASE interface and call the model directly. (i.e. manually convert to AtomicDataDict like calculate and just call the model directly.)

Maybe a model builder can be added that transforms a tensor from irreps to explicit cartesian?

This is a good point. The simplest option of course would be to preprocess your labels into irreps and postprocess your outputs at inference time into Cartesian tensors, but doing this as a module may be easier, depending on the application. At present, nequip also only supports irreps (and not cartesian tensors) in equivariance testing and other model tooling (except for stress/virial tensors) which may also argue in favor of the first solution, unless and until direct support for Cartesian tensors is introduced.