Feat/energy minimizer - #64
Merged
Merged
Conversation
Add minimize_energy (batched torch minimization) and minimize_energy_ase (ASE-driven), an MLCGCalculator ASE wrapper, and unit tests covering both minimizers plus their agreement.
EdoardoRolando
requested review from
aguljas,
jacopoventurin,
kbno and
sayeg84
as code owners
August 5, 2026 17:07
Collaborator
|
This looks great, thanks, although I have a couple of questions:
|
Contributor
|
https://github.com/ClementiGroup/mlcg/blob/feat/ase_calc/src/mlcg/nn/ase_calculator.py Here is a batched calculator FYI. this branch is sort of semi finished. if it is useful, we can talk @EdoardoRolando and combine the two branches? |
Remove the ASE-based MLCGCalculator and minimize_energy_ase from minimizer.py, keeping only the batched, no-ASE minimize_energy. The module no longer imports ASE at all. While auditing minimize_energy itself: - fix a GPU->CPU sync on every structure on every step, from using a CUDA tensor (batch.ptr) directly as a Python slice bound - fix dtype= being silently ignored, so float64 configurations were fed into a float32 model - warn instead of silently returning unconverged geometry when runs out - guard the LBFGS-subclass check against non-class optimizer_cls values (e.g. functools.partial), which issubclass() would reject Expand tests/unit/simulation/test_minimizer.py from 4 to 14 cases: drop the two ASE-only tests and the cross-implementation comparison (nothing left to compare against), add fmax-reached, fixed-atoms- excluded-from-convergence, no-mutation, LBFGS max_iter guard, SGD/Adam support, input validation, converged-structures-stop-being- stepped, mixed-difficulty batches, and the new convergence warning (both that it fires and that it doesn't fire spuriously).
Add mlcg_minimize_energy.py plus parse_minimizer_config() in simulation/cli.py, following the same jsonargparse/--config convention as mlcg_nvt_langevin.py: minimize_energy's own keyword arguments are exposed under a `minimizer:` YAML block via add_function_arguments, alongside model_file/structure_file/output_file. Registers the mlcg-minimize_energy console script and adds an example YAML. Also rename minimize_energy's fixed_atoms parameter to free_atoms, inverting its semantics (list the atoms allowed to move rather than the ones held in place) and keeping None as the "everything moves" default -- shorter and more natural for the common case of relaxing only a handful of atoms. Update the constrained-optimization tests accordingly.
Contributor
Author
|
Thanks Aldo and Zak! I removed the not-bached ase energy minimizer and corresponding mlcg calculator to avoid duplicates, i needed them just as a double check for the pytorch based batched energy minimizer. Thanks @sayeg84 for the suggestion, I have also added a script to interface the energy minimizer from CLI, aside from being usable as a function. |
sayeg84
approved these changes
Aug 18, 2026
sayeg84
left a comment
sayeg84
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, it seems to be good and working.
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
PR Checklist
Adds utilities to relax coarse-grained configurations to a local energy minimum of a trained mlcg model.
What's added (src/mlcg/simulation/minimizer.py):
minimize_energy: batched minimization. All configurations are collated into a single batch and relaxed together with one model forward call per step, using a torch.optim optimizer (LBFGS by default). Model forces are used directly as the negative energy gradient, so no second backward pass is needed.
minimize_energy_ase: relaxes one structure at a time through an ASE optimizer (FIRE by default)
MLCGCalculator: an ASE Calculator wrapper that evaluates energies/forces from an mlcg model, needed for minimize_energy_ase.
Both support per-configuration fixed atoms (frozen during relaxation) and share the same input validation and convergence criterion (max per-atom force magnitude vs fmax).
Tests (tests/unit/simulation/test_minimizer.py): each minimizer is checked for energy decrease, convergence toward equilibrium, and fixed-atom immobility, plus a test that the two methods converge to the same minimum.
Notes: On a real system the batched path was ~15–18× faster than the sequential ASE path while producing matching geometries.