OptimizableBatch and stress relaxations #718
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…nto calculator_updates
…nto calculator_updates
# Conflicts: # src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py
# Conflicts: # src/fairchem/core/trainers/base_trainer.py
rayg1234
reviewed
Nov 27, 2024
rayg1234
previously approved these changes
Dec 2, 2024
zulissimeta
previously approved these changes
Dec 3, 2024
There was a problem hiding this comment.
Approving with trivial update to sphinx autoapi and @rayg1234 's review of everything else!
|
|
||
| # system level model predictions have different shapes than expected by ASE | ||
| ASE_PROP_RESHAPE = MappingProxyType( | ||
| {"stress": (-1, 3, 3), "dielectric_tensor": (-1, 3, 3)} |
There was a problem hiding this comment.
It's a little strange to have these hard-coded. What should a user do implementing new properties?
There was a problem hiding this comment.
I agree this is strange, but it's the simplest solution I thought of. We can always generalized this to simply "tensor of rank X" but in that case we need to have a data structure that provides that information for each model output.
We could use the properties interface and properties defined in ASE to clean this up, but I would suggest doing that in a new PR:
https://gitlab.com/ase/ase/-/blob/master/ase/outputs.py
zulissimeta
approved these changes
Dec 3, 2024
rayg1234
approved these changes
Dec 3, 2024
misko
pushed a commit
that referenced
this pull request
Jan 17, 2025
* remove r_edges, radius, max_neigh and add deprecation warning * edit typing and dont use dicts as default * use super() and remove overkill deprecation warning * set implemented_properties from config * make determine step a method * allow calculator to operate on batches * only update if old config is used * reshape properties * no test classes in ase calculator * yaml load fix * use mappingproxy * expressive import * remove duplicated code * optimizable batch class for ase compatible batch relaxations * fix optimizable batch * optimizable goodies * apply force constraints * use optimizable batch instead and remove torchcalc * update ml relaxations to use optimizable batch correctly * force_consistent check for ASE compat * force_consistent check for ASE compat * check force_consistent * init docs in lbfgs * unitcellfilter for batch relaxations * ruff * UnitCellOptimizable as child class instead of filter * allow running unit cell relaxations * ruff * no grad in run_relaxations * make batched_dot and determine_step methods * imports * rename to optimizableunitcellbatch * allow passing energy and forces explicitly to batch to atoms * check convergence in optimizable and allow passing general results to atoms_from_batch * relaxation test * unit tests * move update mask to optimizable * use energy instead of y * all setting/getting positions and convergence in optimizable * more (unfinished) tests * backwards compatible test * minor fixes * code cleanup * add/fix tests * fix lbfgs * assert using norm * add eps to masked batches if using ASE optimizers * match iterations from previous implementation * use float64 for forces * float32 * use energy_relaxed instead of y_relaxed * energy_relaxed and more explicit error msg * default to batch_size 1 if not set in config * keep float64 training * rename y_relaxed -> energy_relaxed * rm expcell batch * convenience commit from no_experimental_resolve * use numatoms tensor for cell factor * remove positions tests (wrapping atoms gives different results) * allow wrapping positions in batch to atoms * fix test * wrap_positions in batch_to_atoms * take a2g properties from model * test lbfgs traj writes * remove comments * use model generate graph * fix cell_factor * fix using model in ddp * fix r_edges in OCPcalculator * write initial and final structure if save_full is false * check unique atoms saved in trajectory * tighter tol * update ASE release comment * remove cumulative mask option * remove left over cumulative_mask * fix batching when sids as str * do not try to fetch energy and forces if no explicit results * accept Path objects * clean up setting defaults * expose ml_relax in relaxation * force set r_pbc True * make relax_opt optional * no ema on inference only * define ema none to avoid issues * lower force threshold to make sure test does not converge * clean up exception msg * allow strings in batch * remove device argument from lbfgs * minor cleanup * fix optimizable import * do not pass device in ml_relax * simplify enforce max neighbors * fix tests (still not testing stress) * pin sphinx autoapi * typo in version --------- Co-authored-by: zulissimeta <[email protected]> Co-authored-by: Zack Ulissi <[email protected]> Former-commit-id: dc4e598eef45c70e46549e00857968a40f784602
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add functionality to include cell relaxations in structural relaxations for models that implement stress prediction. This PR also enables batch calculations with ASE optimizers.
core.common.relaxation.ml_relaxto enable batch relaxations including stress.OptimizableBatchthat can be used with ASE optimizers or with the existing LBFGS implementationUnitCellOptimizableBatchto allow batch stress relaxationsIncludes awill do this in a separate PRExpCellOptimizableBatchto allow batch stress relaxationsTODO