How to train with custom loss function and different symmetry

[smliu1997]

Hi,

I want to try nequip for building some biomolecular models. So I mainly have two questions:

(1) How to define a custom loss function? For examples, with given sample, I want to compute certain physical variables (such as the gradient of energy to certain variables, and even the Hessian matrix of energy to certain variables) and define the loss function with it, how to do this?
(2) How to change symmetry to SE(3) instead of E(3)? I do not want to keep reflection symmetry.

The tutorials provided perform training with nequip-train, which is convenient but less convenient to extend. If python API is provided, perhaps it is helpful to provide python API tutorials so that we can easily extend the loss functions.

[Linux-cpp-lisp]

Hi @smliu1997 ,

Thanks for your interest in our code!

(1) See https://github.com/mir-group/nequip-example-extension/blob/main/nequip_example_extension/loss_terms.py and that whole repo for custom loss functions. It sounds like you are also asking about predicting custom derivatives: see, for example, https://github.com/mir-group/nequip/blob/main/nequip/nn/_grad_output.py#L115 which gives an idea of how to use other torch.autograd functions.

(2) Depends on what you want---you want predictions that are mirror-equivariant instead of invariant? (i.e. odd parity?)

nequip-train is the only supported interface for model training, but it is fully extensible through callbacks, custom loss functions, custom fields, etc. Workflows built around nequip-train should generally call it as a subprocess like a manual user would.

[smliu1997]

Hi,

Thank you so much for your reply! So basically I want to train a model that does not have reflection symmetry. For example, if the energy of configuration x is E(x), and r(x) is the configuration reflected relative to a mirror, then E(r(x)) can be different from E(x) (basically just SE3 equivariance instead of E3 equivariance). Is this feasible?

Again, thank you so much!