Skip to content

Fixing NNPInput to allow us to write torch script models#295

Merged
chrisiacovella merged 5 commits intochoderalab:mainfrom
chrisiacovella:openmm_integration
Oct 22, 2024
Merged

Fixing NNPInput to allow us to write torch script models#295
chrisiacovella merged 5 commits intochoderalab:mainfrom
chrisiacovella:openmm_integration

Conversation

@chrisiacovella
Copy link
Copy Markdown
Member

@chrisiacovella chrisiacovella commented Oct 22, 2024

Pull Request Summary

When changing NNPInput from a tuple to using slots, pair_list and per atom charge were set to None. While we can call torch.jit.script on the potential with None in the NNPInput and evaluate it, we can't write it to file as it goes through all the validation.

This defines these as empty tensors; we can check the shape of the pair_list in the Neighborlist to know if we need to generate these pairs.

Associated Issue(s)

Pull Request Checklist

  • Issue(s) raised/addressed and linked
  • Includes appropriate unit test(s)
  • Appropriate docstring(s) added/updated
  • Appropriate .rst doc file(s) added/updated
  • PR is ready for review

…as imported into dataset.__init__.py and then being imported from dataset in several place which could be confusing)
…NNPInput. This revises NNPInput and changes the logic used to check if, e.g., pair_list data has been precomputed in the neighborlist for training.
@codecov-commenter
Copy link
Copy Markdown

codecov-commenter commented Oct 22, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 84.97%. Comparing base (88380ce) to head (9b51137).

Additional details and impacted files

Copy link
Copy Markdown
Member

@wiederm wiederm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are good changes! Can you please update the test to cover all implemented potentials?

Comment thread modelforge/jax.py Outdated
Comment thread modelforge/tests/conftest.py Outdated
Comment thread modelforge/tests/test_potentials.py Outdated
Comment thread modelforge/utils/prop.py Outdated
@chrisiacovella chrisiacovella merged commit d6be993 into choderalab:main Oct 22, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants