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

Switch to AEV based models #26

Merged
merged 139 commits into from
Oct 10, 2024
Merged

Switch to AEV based models #26

merged 139 commits into from
Oct 10, 2024

Conversation

lohedges
Copy link
Contributor

This PR updates the code to use the new AEV based EMLE models. The old SOAP based models, i.e. those used for our paper, can be used via the v1.0.0 release of the code. The main changes are as follows:

  • Introduction of fully Torch script compliant models via the emle.models module.
  • Specialised ANI2xEMLE model for one-shot computation of in vacuo and embedding energies (re-using the same AEV features)
  • Introduction of species and reference options for calculation of atomic polaraizabilities.
  • Decoupling of model files from emle-engine repository. Models are now available via emle-models which is managed automatically at run time via git (using the pygit2 package).

kzinovjev and others added 30 commits June 20, 2024 19:04
@kzinovjev
Add emle_features argument to EMLECalculator
b6a434e
@kzinovjev
Create and store aev_computer from ANI-2x when AEV features are used
b940100
@kzinovjev
Proof of concept replacing SOAP with AEV
7ffbeb5
@kzinovjev
Write 'emle-features' value to emle_settings.yaml
e0a297e
@kzinovjev
Keep SOAP features as default for now
1061879
@lohedges
Merge remote-tracking branch 'kzinovjev/soap-to-aev' into feature_aev
c521ccd
@lohedges
Blacken.
c664428
@lohedges
Valiate 'features' kwarg and apply minor formatting tweaks.
f62a107
@lohedges
Expose 'features' in CLI and environment variable.
7942e5e
@lohedges
Update kwarg in error message.
cd630bd
@lohedges
Handle NoneType feature kwarg from emle-server.
e32fa42
@lohedges
Formatting tweaks. [ci skip]
9fa510d
@lohedges
Add support for NNPOps and re-use existing AEVComputer.
7d3e5e5
@lohedges
Switch to using autograd.grad directly.
316c70f
@lohedges
Generalise feature naming.
6d37a76
@lohedges
Fix gradient tensor use with MM embedding.
9338c60
@lohedges
Refactor features validation. [ci skip]
ca2d4da
@lohedges
Extract AEVs from tuple by name. [ci skip]
b3ba1fd
@lohedges
Generalise GPRCalculator docstring.
9798741
@lohedges
Full AEV implementation with features computed in energy function.
1496246 
This fixes the autograd backpropagation by moving the calculation
of the features and derived properties (s, chi) into the energy
function. There is now no need to manually compute feature derivatives
and apply their contribution to the QM gradient via the chain rule.
[ci skip]
@lohedges
AEVCalculator doesn't need device argument. [ci skip]
59d11a0
@lohedges
Remove comment about pt_main_thread issue. [ci skip]
e5346b4
@lohedges
Remove redundant global. [ci skip]
48d83eb
@lohedges
Docstring tweaks. [ci skip]
47edb09
@lohedges
Add initial AEV based EMLE torch.nn.Module model. [ci skip]
82502e0
@lohedges
Call base class constructor to ensure all attributes are initialised.
9a7bc1a 
[ci skip]
@lohedges
Add combined ANI2x-EMLE model. [ci skip]
bb7799e
@lohedges
Rename models module. [ci skip]
80f3e35
@lohedges
Handle situations where there are no point charges.
e16a348
@lohedges
Use same unit system as ANI2x, i.e. Angstrom and Hartree.
74d5f38
@lohedges
Blacken.
db9df87
@lohedges
Refactor to remove code duplication.
4c09678 
Note that TorchScript doesn't support inheritance so we can't derive
from the base EMLE module and call its forward method. Intead, derived
models store instances of both the EMLE and in vacuo models as
attriburtes, then use these to perform the two parts of the calculation.
This avoids re-using the same EMLE logic in each modules forward method.
@lohedges
Clarify setting of AEVs. [ci skip]
6da0239
@lohedges
Clarify pre-computed AEVs. [ci skip]
2184b01
@lohedges
Fix missing _has_nnpops variable.
8e5ffc7
@lohedges
Fix typo in passing device when loading custom MACE models.
627f35a
@lohedges
Move _get_neigbor_pairs back to MACE module now inheritance removed.
15c2b89
@lohedges
Move _get_neigbor_pairs to utility module.
3496c50
@lohedges
Formatting tweak. [ci skip]
1cb7c18
@lohedges
Formatting tweaks. [ci skip]
595a725
@lohedges
Validate the mace_model kwarg.
6f4702f
@lohedges
Miniforge now defaults to mamba. [ci skip]
095fac9
@lohedges
Remove redundant import.
e1835c9
@lohedges
Add basic tests for model instantiation.
8f6dcae
@lohedges
Test that each model can be converted to TorchScript.
1ee46f9
@lohedges
Need to use pip to install test dependencies.
ab38a0c
@lohedges
Need to install using conda and pip.
e74ce0e
@lohedges
Test that models can calculate energies and gradients.
1e7281f
@lohedges
Update test docstrings. [ci skip]
143cb72
@lohedges
Update environment files.
8be242a
@lohedges
Pass user defined model through to ANI2xEMLE constructor.
b046f9f
@lohedges
Update default model names.
70db930
@lohedges
Merge branch 'main' into feature_aev
8514031
@lohedges
Update links for emle-sire tutorial.
9038dfc
@lohedges
Add section on emle.models module.
1a1e9dc
@lohedges lohedges added the enhancement New feature or request label Oct 10, 2024
@lohedges lohedges merged commit 9b8d7c3 into main Oct 10, 2024
5 checks passed
@lohedges lohedges deleted the feature_aev branch October 10, 2024 15:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@kzinovjev kzinovjev kzinovjev approved these changes

Assignees
No one assigned
Labels
enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@lohedges @kzinovjev @JMorado