Introduce ASE-style FIRE optimizer (departing from velocity Verlet in orig FIRE paper) and improve coverage in test_optimizers.py
#174
Conversation
torch_sim/optimizers.py
Outdated
| optimizing atomic positions in a batched setting. Uses velocity Verlet | ||
| integration with adaptive velocity mixing. | ||
| optimizing atomic positions in a batched setting. Logic adapted to follow | ||
| ASE FIRE implementation more closely. |
There was a problem hiding this comment.
I believe that @abhijeetgangan chose to follow the literature implementation intentionally. We've been talking about this internally should offer an implementation that matches ASE but there's a question about how best to do that without duplicating code.
There was a problem hiding this comment.
Understood! In my tests with NEB, I found the ASE FIRE implementation to converge quicker for some reason though. Maybe it's unique to NEB, I can test FireState vs this ASEFireState on a few other structures?
There was a problem hiding this comment.
I can test
FireStatevs thisASEFireStateon a few other structures?
that would be much appreciated!
|
it appears with the changes in this PR, the relaxation @AdeeshKolluru tested in #146 still takes longer to reduce max force than ASE: goes below 1 eV/A in step 17 instead of step 3 for ASE |
…feature/neb-workflow
CompRhys
force-pushed
the
feature/fire-optimizer-changes
branch
from
7f7340c to
25b53aa
Compare
torch_sim/optimizers.py
Outdated
| # ) | ||
| # Get atom-wise masks | ||
| downhill_mask_atoms = downhill_mask_batch[state.batch] | ||
| _uphill_mask_atoms = uphill_mask_batch[state.batch] |
torch_sim/optimizers.py
Outdated
| state: FireState, | ||
| alpha_start: float = alpha_start, | ||
| dt_start: float = dt_start, | ||
| dt_start: float = dt_start, # noqa: ARG001 |
|
Would providing an ASEFireState class be a possible solution? Depending on what the user wants, they can select fire or ase-fire as their optimizer? I have a refactored version of. the ASE FIRE implementation while leaving the closer to paper implementation alone that I can push here. In that implementation I resolved the un-used variables and cleaned things up. The ASEFireState I've written seems to converge as fast as ASE in NEB at least as shown here: #90 (comment) |
that's great to hear! but as you suggest, we need benchmarks across multiple systems to gain some confidence that the ASE implementation is consistently faster to converge
i think to minimize API surface and make it easier to see how the ASE and original FIRE implementations differ, adding a keyword to the existing Args:
model (torch.nn.Module): Model that computes energies, forces, and stress
dt_max (float): Maximum allowed timestep
dt_start (float): Initial timestep
n_min (int): Minimum steps before timestep increase
f_inc (float): Factor for timestep increase when power is positive
f_dec (float): Factor for timestep decrease when power is negative
alpha_start (float): Initial velocity mixing parameter
f_alpha (float): Factor for mixing parameter decrease
maxstep (float): Maximum distance an atom can move per iteration (default
value is 0.2). Only used when md_flavor='ase'.
+ md_flavor ('velocity_verlet' | 'ase'): The type of molecular dynamics flavor to run. Options are
+ 'velocity_verlet' (default) or 'ase' (mimics ASE's FIRE implementation which
+ differs from the original paper). |
|
Sounds good! Will work on this today/tomorrow. |
…step to vv_fire_step. Allowed for selection of md_flavor
…n examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py
|
I just finished modifying optimizer.py and added a test case / example in torch_sim/examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py I tested pure Si (512 atoms), Fe (125 atoms), and Cu (125 atoms), and then Si, Fe, and Cu (same size) with 2 random vacancies in each structure. Here's the output, seems like the ASE implementation is slightly better overall! \n--- Running optimization with MD Flavor: ase_fire ---
Step 50: Active structures: 3, Total converged: 0/5
Step 70: Converged indices [2]. Total converged: 1/5
Step 80: Converged indices [1]. Total converged: 2/5
Step 100: Converged indices [0]. Total converged: 3/5
Step 100: Active structures: 1, Total converged: 3/5
Step 150: Active structures: 2, Total converged: 3/5
Step 180: Converged indices [4]. Total converged: 4/5
Step 200: Active structures: 1, Total converged: 4/5
Step 250: Active structures: 1, Total converged: 4/5
Step 300: Active structures: 1, Total converged: 4/5
Step 310: Converged indices [3]. Total converged: 5/5
All structures converged or max iterations reached by batcher.
Finished ase_fire in 31.12 seconds.
\n--- Running optimization with MD Flavor: vv_fire ---
Step 50: Active structures: 3, Total converged: 0/5
Step 80: Converged indices [2]. Total converged: 1/5
Step 90: Converged indices [1]. Total converged: 2/5
Step 100: Active structures: 1, Total converged: 2/5
Step 130: Converged indices [0]. Total converged: 3/5
Step 150: Active structures: 2, Total converged: 3/5
Step 200: Active structures: 2, Total converged: 3/5
Step 230: Converged indices [4]. Total converged: 4/5
Step 250: Active structures: 1, Total converged: 4/5
Step 300: Active structures: 1, Total converged: 4/5
Step 350: Active structures: 1, Total converged: 4/5
Step 400: Active structures: 1, Total converged: 4/5
Step 450: Active structures: 1, Total converged: 4/5
Step 470: Converged indices [3]. Total converged: 5/5
All structures converged or max iterations reached by batcher.
Finished vv_fire in 43.49 seconds.
\n--- Comparison ---
Force tolerance: 0.05 eV/Å
Initial energies: ['-1546.829', '-375.742', '-899.794', '-140.563', '-240.803'] eV
Final ASE energies: ['-2705.613', '-497.718', '-1053.351', '-2692.331', '-487.226'] eV
Final VV energies: ['-2705.623', '-497.717', '-1053.351', '-2632.801', '-487.219'] eV
Mean Disp (ASE-VV): ['0.0158', '0.0101', '0.0023', '0.9886', '0.0263'] Å
Convergence steps (ASE FIRE): [100, 80, 70, 310, 180]
Convergence steps (VV FIRE): [130, 90, 80, 470, 230]As an aside, I think the bigger differences in the NEB (See #90 ) are due to my batched NEB implementation assigning each image its own FIRE optimizer parameters. Whereas in the ASE version of fire I had used scalar values of timestep, a (mixing parameter), and Nsteps that was used for the entire trajectory (which I think makes more sense to do). I've used Ruff on the optimizer.py. If you guys think this is fine, I can implement the same md_flavor in the UnitCell and FrechetCell versions of FIRE? Also, P.S: InFlightAutoBatcher is Legendary 🚀 Would it be possible to get your feedback on how I used it in the 7.6_Compare_ASE_to_VV_FIRE.py? |
|
@mstapelberg this is great! 👍 thanks a lot i only managed a quick glance. i'll take a closer look tomorrow. porting the same behavior to |
You bet! Will work to port the same behavior to UnitCell and FrechetCell, while doing the same comparisons. Should I adopt/add-on to the test_optimizers.py for this? I notice a lot of the tests there are much shorter/simpler, and don't require a GPU. First time writing proper test cases like this, so happy to give it a go. |
exactly
usually straightforward to port code that works on GPUs to CPUs for unit tests. common strategy for fast tests is to take whatever code you were developing/debugging with and reduce the number of steps and system sizes to as small as possible while still testing something complex enough to cover many edge cases |
… and FrechetCellFireState to have md_flavor to select vv or ase. ASE currently coverges in 1/3 as long. test cases for all three FIRE schemes added to test_optimizers.py with both md_flavors
|
Apologies for the delay on this. I have updated the optimizers.py module to have md_flavor selection for all three types of FIRE. I added test_cases for The "ase-fire" flavor consistently converges in 1/3 the number of steps. |
…nary states, md_flavor validation, non-positive volume warnings brings optimizers.py test coverage up to 96%
…IRE initialization and batch consistency checks maintains same 96% coverage
…_vv_fire_step function modified by functools.partial for different unit cell optimizations (unit/frechet/bare fire=no cell relax) - more concise and maintainable code
instead of _vv_fire_step
torch_sim/optimizers.py
Outdated
| alpha_start: float = 0.1, | ||
| f_alpha: float = 0.99, | ||
| maxstep: float = 0.2, | ||
| md_flavor: MdFlavor = vv_fire_key, |
There was a problem hiding this comment.
Could make ase the default if it consistently converges faster?
There was a problem hiding this comment.
i was also going to suggest that 👍
There was a problem hiding this comment.
Thank you for your kind words Orion!
torch_sim/optimizers.py
Outdated
| - The "vv_fire" flavor follows the original paper closely, including | ||
| integration with Velocity Verlet steps. | ||
| - The "ase_fire" flavor mimics the implementation in ASE, which differs slightly | ||
| in the update steps and does not explicitly use atomic masses in the | ||
| velocity update step. |
There was a problem hiding this comment.
Could be nice to link to the original paper and ASE implementation
…n of FIRE and a link to the original FIRE paper.
…lavor and set default to ase_fire_step
…ws asymmetry in batched mode, batch 0 stalls
fix RuntimeError: a leaf Variable that requires grad is being used in an in-place operation: 7. Position / cell update state.positions += dr_atom
…longer repro error locally
janosh
changed the title
FIRE optimizer (departing from velocity Verlet in orig FIRE paper) and improve coverage in test_optimizers.py
janosh
added
fix
…in fire|unit_cell_fire|frechet_cell_fire doc strings
There was a problem hiding this comment.
thanks a lot @mstapelberg for the massive amount of work that went into this PR! much appreciated! 👍
and apologies for the slow review
* adds comparative test for fire optimizer with ase * rm a few comments * updates CI for the test * update changelog for v0.2.0 (#147) * update changelog for v0.2.0 * minor modification for PR template * formatting fixes * formatting and typos * remove contributors bc they aren't linked * update test with a harder system * fix torch.device not iterable error in test_torchsim_vs_ase_fire_mace.py * fix: should compare the row_vector cell, clean: fix changelog typo * clean: delete .coverage and newline for pytest command * Introduce ASE-style `FIRE` optimizer (departing from velocity Verlet in orig FIRE paper) and improve coverage in `test_optimizers.py` (#174) * feat(fire-optimizer-changes) Update fire_step in optimizers.py based feature/neb-workflow * reset optimizers.py to main version prior to adding updated changes * (feat:fire-optimizer-changes) - Added ase_fire_step and renamed fire_step to vv_fire_step. Allowed for selection of md_flavor * (feat:fire-optimizer-changes) - lint check on optimizers.py with ruff * (feat:fire-optimizer-changes) - added test cases and example script in examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py * (feat:fire-optimizer-changes) - updated FireState, UnitCellFireState, and FrechetCellFireState to have md_flavor to select vv or ase. ASE currently coverges in 1/3 as long. test cases for all three FIRE schemes added to test_optimizers.py with both md_flavors * ruff auto format * minor refactor of 7.6_Compare_ASE_to_VV_FIRE.py * refactor optimizers.py: define MdFlavor type alias for SSoT on MD flavors * new optimizer tests: FIRE and UnitCellFIRE initialization with dictionary states, md_flavor validation, non-positive volume warnings brings optimizers.py test coverage up to 96% * cleanup test_optimizers.py: parameterize tests for FIRE and UnitCellFIRE initialization and batch consistency checks maintains same 96% coverage * refactor optimizers.py: consolidate vv_fire_step logic into a single _vv_fire_step function modified by functools.partial for different unit cell optimizations (unit/frechet/bare fire=no cell relax) - more concise and maintainable code * same as prev commit but for _ase_fire_step instead of _vv_fire_step * (feat:fire-optimizer-changes) - added references to ASE implementation of FIRE and a link to the original FIRE paper. * (feat:fire-optimizer-changes) switched md_flavor type from str to MdFlavor and set default to ase_fire_step * pytest.mark.xfail frechet_cell_fire with ase_fire flavor, reason: shows asymmetry in batched mode, batch 0 stalls * rename maxstep to max_step for consistent snake_case fix RuntimeError: a leaf Variable that requires grad is being used in an in-place operation: 7. Position / cell update state.positions += dr_atom * unskip frechet_cell_fire in test_optimizer_batch_consistency, can no longer repro error locally * code cleanup * bumpy set-up action to v6, more descriptive CI test names * pin to fairchem_core-1.10.0 in CI * explain differences between vv_fire and ase_fire and link references in fire|unit_cell_fire|frechet_cell_fire doc strings * merge test_torchsim_frechet_cell_fire_vs_ase_mace.py with comparative ASE vs torch-sim test for Frechet Cell FIRE optimizer into test_optimizers.py - move `ase_mace_mpa` and `torchsim_mace_mpa` fixtures into `conftest.py` for wider reuse * redirect MACE_CHECKPOINT_URL to mace_agnesi_small for faster tests * on 2nd thought, keep test_torchsim_frechet_cell_fire_vs_ase_mace in a separate file (thanks @CompRhys) * define MaceUrls StrEnum to avoid breaking tests when "small" checkpoints get redirected in mace-torch --------- Co-authored-by: Orion Cohen <[email protected]> Co-authored-by: Janosh Riebesell <[email protected]> Co-authored-by: Rhys Goodall <[email protected]> Co-authored-by: Myles Stapelberg <[email protected]>
Hi @janosh no worries and sorry for the slow reply on my end (was on travel)! Happy to help with this; and excited to use it for my research as well. Thanks for building this great tool! |
|
if you have time, maybe you'd like to give #199 a try to see if it improves your FIRE relaxation results |
|
over at #199 would be great or in a new issue if you find any more problems |
4 participants
This PR adds a new
ase_fireflavor to the FIRE optimizers, mimicking ASE's implementation, alongside the existingvv_fire(Velocity-Verlet based). Users can now select the desired behavior via themd_flavorargument infire,unit_cell_fire, andfrechet_cell_fire. Breaking: The default changed fromvv_firetoase_fire.Highlights:
ase_firesupportsmax_step(unlike originalvv_fire)md_flavor.examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.pycomparesase_fireandvv_firesetup-uvGitHub Action to v6, cleanuppytestconfig inpyproject.toml