Fix discrepancies between FIRE optimisations in ASE and torch-sim
#203
Conversation
This file contains hidden or 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
--------- Co-authored-by: Janosh Riebesell <[email protected]>
* Initialize velocities to None in the pos-only case, see previous changes to the optimizers. (ensures that the correct `dt` is used) * Change the order of the increase of `n_pos`. Again, this ensures the usage of the correct `dt` compared to ASE
* Remove rescaling of positions when updating cell, it's not relevant * Correctly rescale the positions with respect to the deformation gradient * Consider the `cell_forces` in the convergence when doing cell optimizations
|
We require contributors to sign our Contributor License Agreement (CLA), and we don't have record of your signature. In order for us to review and merge your code, please sign the CLA. |
CompRhys
reviewed
May 28, 2025
| dr_atom = torch.where( | ||
| dr_scaling_atom > max_step, max_step * dr_atom / (dr_scaling_atom + eps), dr_atom | ||
| ) | ||
| state.positions = state.positions + dr_atom | ||
|
|
||
| if is_cell_optimization: |
There was a problem hiding this comment.
this seems strange to me? does this logic not add the cartesian dr to the fractional space positions?
There was a problem hiding this comment.
positions = (
torch.linalg.solve(
cur_deform_grad[state.batch], (state.positions+dr_atom).unsqueeze(-1)
).squeeze(-1)
)
state.positions = torch.bmm(
positions.unsqueeze(1), F_new[state.batch].transpose(-2, -1)
).squeeze(1)why would this not be the right thing to do in terms of trying to match your operations here?
There was a problem hiding this comment.
this seems strange to me? does this logic not add the cartesian dr to the fractional space positions?
No, it doesn't scale between fractional/cartesian coordinates. Based on the comments in ASE, this step is done to consider the positions in the original cell. The dr_atom is applied to those coordinates and the new cell deformation is applied afterwards. This is also briefly mentioned in the documentation of the UnitCellFilter (https://wiki.fysik.dtu.dk/ase/ase/filters.html#ase.filters.UnitCellFilter):
Degrees of freedom are the positions in the original undeformed cell, plus the
deformation tensor (extra 3 “atoms”). This gives forces consistent with numerical
derivatives of the potential energy with respect to the cell degreees of freedom.
why would this not be the right thing to do in terms of trying to match your operations here?
Again, to my understanding, you want to consider the positions in the undeformed cell and apply the dr_atom to those. That's why I apply the linalg.solve first and then add the dr_atom, and not the other way around.
CompRhys
reviewed
May 28, 2025
|
I would have to look into the failing test in more detail. As a first step, I took the Still, the test is failing, so I'd have to understand what exactly is happening there. Nonetheless, based on my previous comment, those deviations might be related to the logic of the test and not necessarily to actual differences in the optimizations. I'll look into it and give an update here. Update: I have to stop for now, but there is indeed a problem with the logic of the tests/the loop logic in |
* Include the `cell_forces` in the convergence check * Fix the number of iterations that are performed. `steps_between_swaps` is set to 1, so the number of iterations is equal to the number of swaps. In the previous version, less iterations would have been performed when reaching the maximum number of swaps. For example, when trying to run 32 steps with `steps_between_swaps=5`, the optimization would have stopped after 30 iterations, i.e., 6 swaps. * Fix `autobatching.py`. The if statement would have been triggered for `max_attempts=0`, which was the case when running one iteration and `steps_between_swaps=5`
* Fix the `None` initialization * Fix the cell update when using `UnitCellFilter`
|
The test is now passing, at least locally. As mentioned earlier, one of the issues was the reinitializations between the different steps that are tested. To fix it, I added an if statement for now that checks if the state is already of the correct type. Moreover, as described in the commit message, the current logic of the autobatching caused a different number of iterations, since the structures are set as converged when the maximum number of swap attempts is reached and not necessarily the maximum number of iterations (which led to different number of steps compared to ASE). |
…batch_indices` - fix stale docstring mentioning is_sum_sq kwarg
…issing cell_forces
…ring error messages
janosh
approved these changes
Jun 3, 2025
janosh
added a commit
that referenced
this pull request
Jun 3, 2025
…m` (#203) * fea: use batched vdot --------- Co-authored-by: Janosh Riebesell <[email protected]> * clean: remove ai slop * clean: further attempts to clean but still not matching PR * fix: dr is vdt rather than fdt * typing: fix typing issue * wip: still not sure where the difference is now * update forces per comment * Fix ASE pos only implementation * Initialize velocities to None in the pos-only case, see previous changes to the optimizers. (ensures that the correct `dt` is used) * Change the order of the increase of `n_pos`. Again, this ensures the usage of the correct `dt` compared to ASE * Fix torch-sim ASE-FIRE (Frechet Cell) * Remove rescaling of positions when updating cell, it's not relevant * Correctly rescale the positions with respect to the deformation gradient * Consider the `cell_forces` in the convergence when doing cell optimizations * linting * test: still differ significantly after step 1 for distorted structures * Fix test comparing ASE and torch-sim optimization * Include the `cell_forces` in the convergence check * Fix the number of iterations that are performed. `steps_between_swaps` is set to 1, so the number of iterations is equal to the number of swaps. In the previous version, less iterations would have been performed when reaching the maximum number of swaps. For example, when trying to run 32 steps with `steps_between_swaps=5`, the optimization would have stopped after 30 iterations, i.e., 6 swaps. * Fix `autobatching.py`. The if statement would have been triggered for `max_attempts=0`, which was the case when running one iteration and `steps_between_swaps=5` * Fix `optimizers` when using `UnitCellFilter` * Fix the `None` initialization * Fix the cell update when using `UnitCellFilter` * fix test_optimize_fire * allow FireState.velocities = None since it's being set to None in multiple places * safer `batched_vdot`: check dimensionality of input tensors `y` and `batch_indices` - fix stale docstring mentioning is_sum_sq kwarg * generate_force_convergence_fn raise informative error on needed but missing cell_forces * pascal case VALID_FIRE_CELL_STATES->AnyFireCellState and fix non-f-string error messages * fix FireState TypeError: non-default argument 'dt' follows default argument * allow None but don't set default for state.velocities * fix bad merge conflict resolution * tweaks --------- Co-authored-by: Rhys Goodall <[email protected]> Co-authored-by: Janosh Riebesell <[email protected]>
|
@t-reents thank you for your careful work here! |
janosh
added a commit
that referenced
this pull request
Jun 9, 2025
* clean: move pair fn logic into batched model code * clean: begin to remove unbatched scripts * clean: update the dynamics scripts, remove old code, refactor the integrators into folder * fix: all unit tests run * fix: remove circular import * wip: address more issues with removing unbatched code * wip: another attempt to fix examples * fix: mace needs the shifts? we should add a unit test that needs them as the MaceModel unit tests didn't fail removing them. * metatensor models have been renamed to metatomic models (#205) * Fix discrepancies between `FIRE` optimisations in `ASE` and `torch-sim` (#203) * fea: use batched vdot --------- Co-authored-by: Janosh Riebesell <[email protected]> * clean: remove ai slop * clean: further attempts to clean but still not matching PR * fix: dr is vdt rather than fdt * typing: fix typing issue * wip: still not sure where the difference is now * update forces per comment * Fix ASE pos only implementation * Initialize velocities to None in the pos-only case, see previous changes to the optimizers. (ensures that the correct `dt` is used) * Change the order of the increase of `n_pos`. Again, this ensures the usage of the correct `dt` compared to ASE * Fix torch-sim ASE-FIRE (Frechet Cell) * Remove rescaling of positions when updating cell, it's not relevant * Correctly rescale the positions with respect to the deformation gradient * Consider the `cell_forces` in the convergence when doing cell optimizations * linting * test: still differ significantly after step 1 for distorted structures * Fix test comparing ASE and torch-sim optimization * Include the `cell_forces` in the convergence check * Fix the number of iterations that are performed. `steps_between_swaps` is set to 1, so the number of iterations is equal to the number of swaps. In the previous version, less iterations would have been performed when reaching the maximum number of swaps. For example, when trying to run 32 steps with `steps_between_swaps=5`, the optimization would have stopped after 30 iterations, i.e., 6 swaps. * Fix `autobatching.py`. The if statement would have been triggered for `max_attempts=0`, which was the case when running one iteration and `steps_between_swaps=5` * Fix `optimizers` when using `UnitCellFilter` * Fix the `None` initialization * Fix the cell update when using `UnitCellFilter` * fix test_optimize_fire * allow FireState.velocities = None since it's being set to None in multiple places * safer `batched_vdot`: check dimensionality of input tensors `y` and `batch_indices` - fix stale docstring mentioning is_sum_sq kwarg * generate_force_convergence_fn raise informative error on needed but missing cell_forces * pascal case VALID_FIRE_CELL_STATES->AnyFireCellState and fix non-f-string error messages * fix FireState TypeError: non-default argument 'dt' follows default argument * allow None but don't set default for state.velocities * fix bad merge conflict resolution * tweaks --------- Co-authored-by: Rhys Goodall <[email protected]> Co-authored-by: Janosh Riebesell <[email protected]> * bump pre-commit hooks and fix new errors * fix MACE examples not using ts.io state conversion utilities * fix unused shifts_list not torch.cat()-ed in mace.py * fix examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py using nvt_langevin in place of nvt_nose_hoover_invariant remove debug print in torch_sim/runners.py * just calc total energy manually in nvt_langevin_invariant in torch_sim/integrators/nvt.py also fix broadcasting bug in Ornstein-Uhlenbeck (ou_step) * fix calculate_momenta() in npt_nose_hoover_init also fix 3.8_MACE_NPT_Nose_Hoover.py broken imports from torch_sim.integrators.npt * fix unpack of n_particles, dim = state.positions.shape * fix missing external_pressure in npt_nose_hoover call in 3.8_MACE_NPT_Nose_Hoover.py * update cell_position in npt_nose_hoover before downstream usage * refactor (no bug fixes) * fix npt_nose_hoover transposing cell that's already in column-vector convention * feat: Add batching support to NPT Nose-Hoover integrator - Update NPTNoseHooverState cell properties to support batch dimensions - reference_cell: [n_batches, 3, 3] instead of [3, 3] - cell_position, cell_momentum, cell_mass: [n_batches] instead of scalar - Fix tensor broadcasting in exp_iL1, exp_iL2 with proper atom-to-batch mapping - Update compute_cell_force for per-batch kinetic energy and stress calculations - Fix npt_nose_hoover_init to properly initialize batched cell variables - Update npt_nose_hoover_invariant for per-batch energy conservation - Replace pbc_wrap_general with pbc_wrap_batched for proper PBC handling - Fix example script 3.8_MACE_NPT_Nose_Hoover.py output formatting Enables multiple independent NPT systems in a single simulation while maintaining backward compatibility for single-batch systems. * fix: Remove cell dimension squeezing in NVT Nose-Hoover integrator - Remove problematic cell.squeeze(0) that breaks batching support - Fix calculate_momenta function call to use correct signature with batch parameter - Resolves RuntimeError when using MACE with NVT Nose-Hoover thermostat Fixes example script 3.5_MACE_NVT_Nose_Hoover.py which was failing due to neighbor list function receiving wrong tensor shapes when cell batch dimension was incorrectly removed. * fix: Handle scalar kT and batched stress tensors in NPT Nose-Hoover - Convert scalar kT to tensor before accessing .ndim attribute in npt_nose_hoover_init and update_cell_mass - Fix stress tensor trace computation in compute_cell_force to handle 3D batched tensors - Use torch.diagonal().sum() for batched stress tensors instead of torch.trace() Fixes Lennard-Jones NPT Nose-Hoover script that was failing with: - AttributeError: 'float' object has no attribute 'ndim' - RuntimeError: trace: expected a matrix, but got tensor with dim 3 * fix: Handle batched cell tensors in get_fractional_coordinates - Replace deprecated .T with .mT for matrix transpose on 3D tensors - Add support for batched cell tensors with shape [n_batches, 3, 3] - Extract first batch cell matrix when cell.ndim == 3 - Maintains backward compatibility with 2D cell matrices Fixes batched silicon workflow script that was failing with: - UserWarning about deprecated .T usage on >2D tensors - RuntimeError: linalg.solve: A must be batches of square matrices The get_fractional_coordinates function now properly handles both single [3,3] and batched [n_batches,3,3] cell tensors, enabling a2c_silicon_batched.py workflow to run * replace .transpose(-2, -1) with .mT everywhere * fix nvt + npt integrators: ensure consistent batch dimensions in kinetic energy calculations - Add missing `batch=state.batch` parameter to `calc_kinetic_energy` calls in NPT integrator - Enhance NVT/NPT invariant functions with explicit broadcasting for chain variables - Replace inefficient manual loops with batch-aware kinetic energy calculation in NPT invariant - Fix undefined variable reference in NPT invariant function (n_batches → state.n_batches) Resolves broadcasting issues where scalar kinetic energy was incorrectly combined with batched energy and temperature tensors, ensuring all energy terms have consistent batch dimensions before addition. * fix transforms.py: prevent silent data corruption in get_fractional_coordinates Replace problematic batched cell handling that only processed the first batch (cell[0]) with explicit NotImplementedError. This prevents silent data corruption where multi-batch systems would incorrectly use only the first batch's cell parameters for all coordinate transformations. - Raise NotImplementedError for 3D batched cell tensors instead of silently ignoring batches 1, 2, ... N - Preserve existing functionality for 2D cell matrices - Add clear error message indicating limitation and suggesting workarounds Breaking change: Code that previously silently failed will now explicitly error, but this prevents incorrect results in multi-batch scenarios. * fix nvt_langevin: handle None gamma parameter correctly * fix npt: missing PBC check in Nose-Hoover position update - Check state.pbc before applying ts.transforms.pbc_wrap_batched - Return unwrapped positions when state.pbc is False - Ensure consistent behavior with NPT Langevin implementation * Initialize cur_deform_grad to prevent UnboundLocalError * fix nvt: compute degrees of freedom per batch in Nose-Hoover init Replace count_dof() with proper batch-aware DOF calculation using torch.bincount(state.batch) to ensure consistency with invariant function. Previously ignored batch structure, causing incorrect DOF for batched systems. * fix /5.1_a2c_silicon_batched.py: restore single-batch cell tensors support in transforms.py's get_fractional_coordinates - Handle batched cell tensors with shape [1, 3, 3] by auto-squeezing to [3, 3] - Improve error messages for multi-batch cases to be more informative - Add comprehensive tests for batched cell tensor scenarios - Fixes NotImplementedError in examples/scripts/5_Workflow/5.1_a2c_silicon_batched.py:159 - Maintains full backward compatibility with existing 2D cell matrix usage --------- Signed-off-by: Janosh Riebesell <[email protected]> Co-authored-by: Janosh Riebesell <[email protected]>
|
Warning Rate limit exceeded@janosh has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 4 minutes and 9 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (8)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
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.
It seems that the remaining deviations could be fixed. This branch is already based on the changes in #201, I'll do the rebase once it is merged.
1, in order to get the precise number of convergence stepsdtgenerate_force_convergence_fnfunction to optionally also include thecell_forcesin the convergence check, which is done inASEand therefore caused the differing numbers of convergence steps