Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Flexible property prediction heads #362

Draft
wants to merge 47 commits into
base: main
Choose a base branch
from
Draft

Flexible property prediction heads #362

wants to merge 47 commits into from

Conversation

M-R-Schaefer
Copy link
Contributor

@M-R-Schaefer M-R-Schaefer commented Nov 4, 2024 •
edited
Loading

I have added property prediction heads to the EnergyModel.
These can predict per-atom or per-structure quantities with either l0, l1 or symmetric traceless l2 symmetry.

Further, I have added a LatentEwald module which gives us access to long range interactions

TODO:

  • add has_aux to all gradient transformations
  • fix jaxmd
  • fix shallow ensemble (how do we make this compatible with e.g. ewald?)
  • add arbitrary property readout to PBP dataset
  • tests

M-R-Schaefer and others added 24 commits October 30, 2024 23:31
@M-R-Schaefer
first sketches
e8e9ef8
@M-R-Schaefer
removed atomistic model, added properties dict
79872b6
@M-R-Schaefer
removed atomisticmodel import
c3c950b
@M-R-Schaefer
first draft of latent ewald sum
9569461
@M-R-Schaefer
working implementation of property heads
2509291
@M-R-Schaefer
added property head to builder
70e0fc2
@M-R-Schaefer
fixed stress and added property shift
14e7f52
@M-R-Schaefer
moved divisor out of inner loss function
6fc2a98
@M-R-Schaefer
fixed properties head
a1b2957
@M-R-Schaefer
removed outdated comment
ce4da9c
@M-R-Schaefer
removed old additional properties flag
b7d35cd
@M-R-Schaefer
added additional properties to dataset
596cd49
@M-R-Schaefer
fixed loss fns, added additinoal prop to convert
232742d
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
897c3af 
for more information, see https://pre-commit.ci
@M-R-Schaefer
fixed bal and shallow ens
501a4bb
@M-R-Schaefer
remove atomistic model
36ca219
@M-R-Schaefer
fixed template
af2aaf2
@M-R-Schaefer
fixed jaxmd and shallow ens compat
f339cdf
@M-R-Schaefer
apax nodes eager mode
f623669
@M-R-Schaefer
no longer error when restarting finished training
bf8808a
@M-R-Schaefer
fixed remaining tests
0a73502
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
14f1e89 
for more information, see https://pre-commit.ci
@M-R-Schaefer
Merge branch 'main' into dtensor
64f77a3
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
0a568e5 
for more information, see https://pre-commit.ci
pre-commit-ci bot and others added 14 commits November 17, 2024 13:16
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
945c745 
for more information, see https://pre-commit.ci
@M-R-Schaefer
moved dtype parsing into models
2bc9f57
@M-R-Schaefer
more dtype refactoring
29a4554
@M-R-Schaefer
atom mask compatibility with arbitrary sized arrays
6a7b28a
@M-R-Schaefer
removed old additional proeprteis docstring
b6681fd
@M-R-Schaefer
remove debug comments
bb7359f
@M-R-Schaefer
detect which labels to get via loss config
073aba8
@M-R-Schaefer
Merge branch 'dtensor' of https://github.com/apax-hub/apax into dtensor
9fdd618
@M-R-Schaefer
Merge branch 'main' into dtensor
bda5a04
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
9bd3075 
for more information, see https://pre-commit.ci
@M-R-Schaefer
fixed tests
f6705f2
@M-R-Schaefer
Merge branch 'dtensor' of https://github.com/apax-hub/apax into dtensor
4f9f58f
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
25c6e31 
for more information, see https://pre-commit.ci
@M-R-Schaefer
Merge branch 'main' into dtensor
522494d
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@M-R-Schaefer