torchmd-net icon indicating copy to clipboard operation
torchmd-net copied to clipboard

Published 20 hours ago • verified publisher icon torchmd

Support atom-wise and molecule-wise priors

Open raimis opened this issue 2 years ago • 1 comments

Implement the support of atom-wise (e.g. atomic reference energies) and molecule-wise (e.g. electrostatic term) priors.

  • [x] Refactor the priors
  • [x] Implement pre-reduce and post-reduce methods

raimis avatar Oct 25 '22 13:10 raimis

This conflicts with #134, which implements the same feature in a different way. Can we agree on which API we want to use? I followed what @PhilippThoelke suggested. If we decide to use the other API, let me change it there so we don't create conflicts.

peastman avatar Oct 25 '22 16:10 peastman

Key differences with #134:

  • The atom- and molecular-wise operations are implement with pre_reduce and post_reduce methods. This is consistent with the output model.
  • The priors are executed after the corresponding output models
  • The molecule-wise prior signatures is: 
    y = self.prior_model.post_reduce(y, z, pos, batch)
    This more flexible than just the additive term.

The other things can be as proposed by @peastman

raimis avatar Oct 26 '22 14:10 raimis

@stefdoerr @PhilippThoelke what do you think?

raimis avatar Oct 26 '22 14:10 raimis

I think both #134 and this PR will be useful and as far as I understand it they're not mutually exclusive? @peastman or is the conflict more substantial? I like the idea of pre_ and post_reduce but that should still work for the ZBL prior and the changes you made to the args?

PhilippThoelke avatar Oct 26 '22 14:10 PhilippThoelke

I think both https://github.com/torchmd/torchmd-net/pull/134 and this PR will be useful and as far as I understand it they're not mutually exclusive?

We just need to pick which API we prefer for the priors. Either one can work for my needs. I don't think it would make sense to do both, since they're basically alternate APIs for doing the same thing.

peastman avatar Oct 26 '22 16:10 peastman

I would say the API in this PR is more general and therefore preferable but I think it's still useful to keep e.g. the distinction between prior_args and prior_init_args from #134.

PhilippThoelke avatar Oct 27 '22 03:10 PhilippThoelke

I agree with Philipp, this implementation seems more generic but should be trivial to merge with the additional changes in #134

stefdoerr avatar Oct 27 '22 07:10 stefdoerr

So, the consensus is to use the API from this PR and the argument part from #134. Could somebody approve and I'll merge?

raimis avatar Oct 27 '22 15:10 raimis