GROMACS engine support to a3fe - #76
Roy-Haolin-Du wants to merge 12 commits into
Conversation
|
Hi @fjclark , this PR #76 is now ready for review. It adds non-adaptive GROMACS ABFE support, including charged ligands, analysis, regression tests, documentation, and a minimal SLURM integration test. I also validated the implementation (gmx engine) locally using four ligands (two neutral and two charged, 0.1ns); So, when you have a chance, would you mind taking a look? Best wishes, |
| _StageType.VANISH: [ | ||
| 0.0, | ||
| 0.1, | ||
| 0.2, | ||
| 0.3, | ||
| 0.4, | ||
| 0.5, | ||
| 0.55, | ||
| 0.6, | ||
| 0.65, | ||
| 0.7, | ||
| 0.75, | ||
| 0.8, | ||
| 0.85, | ||
| 0.9, | ||
| 0.95, | ||
| 1.0, | ||
| ], |
There was a problem hiding this comment.
More science than software question, but it would be nice to document (even if that's just adding to the description) where the lambda values came from (my fault for not doing this originally).
The most consistent way to do this would be to space the default windows with the same thermodynamic speed for the same system for GROMACs and SOMD (but feel free to leave for now if that's a pain).
| ligand_charge: int = _Field( | ||
| 0, | ||
| description="Net charge of the ligand. If non-zero, must use PME for electrostatics.", | ||
| description="Charge change for the alchemical transformation. If non-zero, PME must be used.", |
There was a problem hiding this comment.
Is "Net charge of the ligand" less ambiguous? It's not defined which direction the charge change is in so e.g. could be plus or minus the ligand charge.
| pcoupl: _Literal["no", "Berendsen", "C-rescale", "Parrinello-Rahman"] = _Field( | ||
| "Parrinello-Rahman", description="Pressure coupling" | ||
| ) |
There was a problem hiding this comment.
From a quick skim your defaults look good to me! If you've taken them from anywhere in particular e.g. BioSimSpace or https://github.com/bigginlab/fragment-opt-abfe-benchmark/blob/main/simulation_control_files/abfe_mdps/4fs/coul/prod/prod.mdp, would be good to note this in the docstring.
| if self.sims[0].engine_type == _EngineType.GROMACS: | ||
| # GROMACS: dt in ps, nstdhdl is steps (dH/dlambda output frequency), convert to ns | ||
| # First energy is written at time 0, so no offset needed | ||
| time_per_energy = config.dt * config.nstdhdl / 1000 | ||
| equil_index = int(self._equil_time / time_per_energy) | ||
| else: | ||
| # SOMD: timestep in fs, energy_frequency is steps, convert to ns | ||
| # First energy is only written after the first nrg_freq steps, so subtract 1 | ||
| time_per_energy = config.timestep * config.energy_frequency / 1_000_000 | ||
| equil_index = max(0, int(self._equil_time / time_per_energy) - 1) |
There was a problem hiding this comment.
I think it would be neater to have this logic inside each engine config, so that Lambda window (and as far as possible, everything in run) doesn't have to know anything about the engines -- it just calls e.g. equil_index = config.get_equil_index(self._equil_time).
| if sim.engine_type == _EngineType.GROMACS: | ||
| in_file = sim.output_dir + "/prod/prod.xvg" | ||
| out_file = sim.output_dir + "/prod/prod_equilibrated.xvg" | ||
| header_chars = ("#", "@") | ||
| else: | ||
| in_file = sim.output_dir + "/simfile.dat" | ||
| out_file = sim.output_dir + "/simfile_equilibrated.dat" | ||
| header_chars = ("#",) |
There was a problem hiding this comment.
As above, I think it would be neater if engine-specific paths were handled by an engine-specific class so that LamWindow can remain completely blind to what engine is used -- it would be great to avoid engine-specific if else statements in run.
| lines = ifile.readlines() | ||
|
|
||
| # Figure out how many lines come before the data | ||
| non_data_lines = 0 | ||
| for line in lines: | ||
| if line.startswith("#"): | ||
| if line.startswith(header_chars) or not line.strip(): |
There was a problem hiding this comment.
If this is engine-specific, again it would be cool if were in a class specific to an engine.
| @@ -452,7 +452,7 @@ def run_ensemble_equilibration( | |||
| print( | |||
| f"Running ensemble equilibration simulation with GROMACS for {cfg.ensemble_equilibration_time} ps" | |||
| ) | |||
| if leg_type == _LegType.BOUND: | |||
| if leg_type == _LegType.BOUND or engine_type == _EngineType.GROMACS: | |||
There was a problem hiding this comment.
Here and below in this file (and in run in general), it would be really nice and clean if we could move all of the engine-specific logic out of run, as far as possible (if you think it's too much of a pain to be worth it, feel free to ignore). It feels like this decision might be best made in the engine-specific system prep config.
| def _add_gromacs_alchemical_ions( | ||
| system: _BSS._SireWrappers._system.System, # type: ignore | ||
| ligand: _BSS._SireWrappers._molecule.Molecule, # type: ignore | ||
| ligand_charge: int, | ||
| ) -> None: | ||
| """Make existing counterions co-alchemical for a GROMACS decoupling.""" | ||
| ion_charge = -1 if ligand_charge > 0 else 1 | ||
| ions = [ | ||
| mol | ||
| for mol in system | ||
| if mol.nAtoms() == 1 and round(mol.charge().value()) == ion_charge | ||
| ] | ||
| if len(ions) < abs(ligand_charge): | ||
| raise ValueError( | ||
| f"Could not find {abs(ligand_charge)} monovalent counterion(s) " | ||
| "for the charged ligand." | ||
| ) | ||
|
|
||
| space = system._sire_object.property("space") | ||
| ligand_centre = ligand.getAtoms()[ligand.getCOMIdx()]._sire_object.property( | ||
| "coordinates" | ||
| ) | ||
| ions.sort( | ||
| key=lambda ion: space.calc_dist( | ||
| ion.getAtoms()[0]._sire_object.property("coordinates"), ligand_centre | ||
| ), | ||
| reverse=True, | ||
| ) | ||
|
|
||
| for ion in ions[: abs(ligand_charge)]: | ||
| perturbed_ion = _BSS.Align.merge(ion, ion, mapping={0: 0}) | ||
| cursor = perturbed_ion._sire_object.cursor() | ||
| charge = perturbed_ion.getAtoms()[0]._sire_object.property("charge1") | ||
| cursor[0]["charge1"] = 0 * charge | ||
| perturbed_ion._sire_object = cursor.commit() | ||
| system.updateMolecule(system.getIndex(ion), _mark_alchemical_ion(perturbed_ion)) |
There was a problem hiding this comment.
Saying the same thing again, but would be neat to move all the engine-specific stuff out of run, or at least into something like run/gromacs_helpers.py
| def _ensemble_equilibration_output_files(self) -> _List[str]: | ||
| """Return the files expected from ensemble equilibration.""" | ||
| if self.engine_type == _EngineType.GROMACS: | ||
| files = ["gromacs.gro"] | ||
| if self.leg_type == _LegType.BOUND: | ||
| files.append("gromacs.xtc") | ||
| return files | ||
| return ["somd.rst7"] |
There was a problem hiding this comment.
Would be cool to move this logic inside the EngineType classes so that all the logic in run is engine-independent. This is dependency inversion -- we want to depend on the interfaces (e.g. knowing that EngineType has a get_files function) rather than implementing engine-level logic here. This would make it easier in future to add another engine (not that we're planning that, but it keeps the code nice and separates concerns).
| coordinate_prefix, coordinate_extension = ( | ||
| ("gromacs", "gro") | ||
| if self.engine_type == _EngineType.GROMACS | ||
| else ("somd", "rst7") | ||
| ) |
There was a problem hiding this comment.
I would make coordinate_prefix_and_extension a property of EngineType.
| @@ -603,7 +656,10 @@ def run_ensemble_equilibration( | |||
|
|
|||
| # Save the restraints to a text file and store within the Leg object | |||
| with open(f"{outdir}/restraint_{i + 1}.txt", "w") as f: | |||
| f.write(restraint.toString(engine="SOMD")) # type: ignore | |||
| run_engine = ( | |||
| "GROMACS" if self.engine_type == _EngineType.GROMACS else "SOMD" | |||
There was a problem hiding this comment.
I would again move this logic inside EngineType so you can go e.g.
run_engine = self.engine_type.run_engine and move out engine-specific logic.
| for file in _glob.glob(f"{stage_input_dir}/lambda_0.0000/*") | ||
| if not file.endswith(".cfg") | ||
| for file in _glob.glob(f"{stage_input_dir}/lambda_*/*") | ||
| if not file.endswith((".cfg", ".mdp", ".tpr")) |
There was a problem hiding this comment.
It would be nice to collect these somewhere else, e.g. in EngineType.config_file_suffixes. Then you could combine them all into a constant in configuration do from ..configuration import CONFIG_FILE_SUFFIXES as _CONFIG_FILE_SUFFIXES, and here to `if not file.endswith(_CONFIG_FILE_SUFFIXES).
| self.engine_config.ref_t | ||
| if self.engine_type == _EngineType.GROMACS | ||
| else 298.15 |
There was a problem hiding this comment.
Would be great if this logic was in engine-specific modules/ classes
| required_input_files = { | ||
| _EngineType.SOMD: [ | ||
| "somd.prm7", | ||
| "somd.rst7", | ||
| "somd.pert", | ||
| ], | ||
| _EngineType.GROMACS: [ | ||
| "gromacs.top", | ||
| "gromacs.gro", | ||
| ], | ||
| } |
There was a problem hiding this comment.
Might be neater if these were a property of EngineType so you could do e.g. engine_type.required_input_files
| if self.engine_type == _EngineType.GROMACS: | ||
| # Move GROMACS checkpoint files | ||
| for cpt_file in _glob.glob( | ||
| f"{self.output_dir}/**/*.cpt", recursive=True | ||
| ): | ||
| _subprocess.run(["mv", cpt_file, f"{self.output_dir}/failure"]) | ||
| else: # SOMD |
There was a problem hiding this comment.
Again would be great to avoid engine-specific behaviour in run
| """Select the coordinate and restraint files for this run.""" | ||
| if self.engine_type == _EngineType.SOMD: | ||
| file_prefix = "somd" | ||
| file_extension = "rst7" | ||
| elif self.engine_type == _EngineType.GROMACS: | ||
| file_prefix = "gromacs" | ||
| file_extension = "gro" |
There was a problem hiding this comment.
Would be great for this to go into an engine-specific class or module
| assert results_exp.loc["t4l", "calc_dg"] == pytest.approx(5.2622, abs=1e-2) | ||
| assert results_exp.loc["t4l", "calc_er"] == pytest.approx(0.1138, abs=1e-2) | ||
| assert results_exp.loc["t4l", "calc_dg"] == pytest.approx(5.2850, abs=1e-2) | ||
| assert results_exp.loc["t4l", "calc_er"] == pytest.approx(0.0808, abs=1e-2) |
There was a problem hiding this comment.
Any idea why the regression test has broken? Was there a bug in the original?
There was a problem hiding this comment.
Yes, this was caused by a bug in the original equilibration-index calculation. equil_time is in ns while timestep is in fs.
This was fixed
a3fe/a3fe/run/lambda_window.py
Lines 354 to 363 in 4d7b9a7
main.Thanks!
| mbar_temperature = ( | ||
| lambda_windows[0].engine_config.ref_t | ||
| if engine_type == _EngineType.GROMACS | ||
| else 298.15 | ||
| ) |
There was a problem hiding this comment.
Would be nice to move engine-specific logic into engine-specific classes or modules
| def _run_mbar_gromacs( | ||
| output_dir: str, | ||
| run_nos: _List[int], | ||
| percentage_end: float = 100, | ||
| percentage_start: float = 0, | ||
| subsampling: bool = False, | ||
| delete_outfiles: bool = False, | ||
| equilibrated: bool = True, | ||
| temperature: float = 298.15, | ||
| ) -> _Tuple[_np.ndarray, _np.ndarray, _List[str], _Dict[str, _Dict[str, _np.ndarray]]]: | ||
| """ | ||
| Run MBAR on GROMACS dhdl output files using alchemlyb and pymbar. | ||
| """ | ||
| if subsampling: | ||
| raise NotImplementedError("Subsampling is not implemented for GROMACS MBAR.") | ||
|
|
||
| import pandas as _pd | ||
| import pymbar as _pymbar | ||
| from alchemlyb.parsing.gmx import extract_u_nk as _extract_u_nk | ||
|
|
||
| tmp_files = _prepare_xvg_files( | ||
| output_dir=output_dir, | ||
| run_nos=run_nos, | ||
| percentage_end=percentage_end, | ||
| percentage_start=percentage_start, | ||
| equilibrated=equilibrated, | ||
| ) | ||
|
|
||
| mbar_out_files = [] | ||
| free_energies = [] | ||
| errors = [] | ||
| mbar_grads = {} | ||
| kt = _R * temperature / 4184 | ||
|
|
||
| for run_no in run_nos: | ||
| xvg_files = _get_gromacs_xvg_files( | ||
| output_dir=output_dir, | ||
| run_no=run_no, | ||
| percentage_end=percentage_end, | ||
| percentage_start=percentage_start, | ||
| ) | ||
| u_nk = _pd.concat( | ||
| [_extract_u_nk(xvg_file, temperature) for xvg_file in xvg_files] | ||
| ) | ||
| u_nk = u_nk.sort_index(level=u_nk.index.names[1:]) | ||
| groups = u_nk.groupby(level=u_nk.index.names[1:]) | ||
| n_k = [ | ||
| len(groups.get_group(state)) if state in groups.groups else 0 | ||
| for state in u_nk.columns | ||
| ] | ||
| mbar = _pymbar.MBAR(u_nk.T, n_k, initialize="BAR") | ||
| delta_f, d_delta_f, _ = mbar.getFreeEnergyDifferences(return_theta=True) | ||
|
|
||
| outfile = ( | ||
| f"{output_dir}/freenrg-MBAR-run_{str(run_no).zfill(2)}_" | ||
| f"{round(percentage_end, 3)}_end_" | ||
| f"{round(percentage_start, 3)}_start.dat" | ||
| ) | ||
| mbar_out_files.append(outfile) | ||
|
|
||
| lam_vals = _get_lambda_values_from_xvg_files(xvg_files) | ||
| delta_f = _np.array(delta_f) * kt | ||
| d_delta_f = _np.array(d_delta_f) * kt | ||
| overlap = _get_overlap_matrix(mbar) | ||
|
|
||
| free_energies.append(delta_f[0, -1]) | ||
| errors.append(d_delta_f[0, -1]) | ||
|
|
||
| _write_gromacs_mbar_output( | ||
| outfile=outfile, | ||
| xvg_files=xvg_files, | ||
| lam_vals=lam_vals, | ||
| delta_f=delta_f, | ||
| d_delta_f=d_delta_f, | ||
| overlap=overlap, | ||
| subsampling=subsampling, | ||
| ) | ||
|
|
||
| mbar_grads[f"run_{run_no}"] = {} | ||
| lam_vals, grads, grad_errors = _read_mbar_gradients(outfile) | ||
| mbar_grads[f"run_{run_no}"]["lam_vals"] = lam_vals | ||
| mbar_grads[f"run_{run_no}"]["grads"] = grads | ||
| mbar_grads[f"run_{run_no}"]["grad_errors"] = grad_errors | ||
|
|
||
| if delete_outfiles: | ||
| for ofile in mbar_out_files: | ||
| _subprocess.run(["rm", ofile]) | ||
| mbar_out_files = [] | ||
|
|
||
| for tmp_file in tmp_files: | ||
| _subprocess.run(["rm", tmp_file]) | ||
|
|
||
| return _np.array(free_energies), _np.array(errors), mbar_out_files, mbar_grads |
There was a problem hiding this comment.
Would be great to test this against MBAR run through e.g. BioSimSpace (if possible) or even just checking it's fairly close to the GROMACS BAR result.
There was a problem hiding this comment.
Hi @Roy-Haolin-Du, thanks and congrats! Great to see this functionality.
I've had a skim over everything and looked at most of run in more detail.
Overall this looks really good! My main comment (reflected in most of my smaller comments) is that it would be amazing if we could separate out all of the engine specific logic as much as possible from other parts of a3fe, for example it would be great if run was almost completely engine-agnostic. In practice, this mainly mains eliminating engine if else statements like "if gromacs do this, else if somd do this", moving the logic into an engine specific module/ function so that the run modules can just depend on an interface rather than implementing engine-specific logic themselves. In practice it might not be practical to 100 % extract this logic but I think even moving towards this would make things a bit cleaner.
Also, I've started getting AI to review all my PRs recently and have found this often catches bugs. As well as the standard /review skills, I've found /validate-science fromhttps://github.com/dotsdl/agents-for-scientific-software useful, and also often use https://github.com/DietrichGebert/ponytail to review for conciseness and suggest where code could be removed.
I've not got time to check these (so ignore if irrelevant) but here's the result of asking codex to review and validate science (mainly issues around statistical inefficiency so not a massive problem since we're using inter-replicate errors, but obviously still worth addressing if correct):
Review:
- [P1] Clean nested GROMACS outputs before rerunning — a3fe/run/simulation.py:310
SimulationRunner.clean() only performs non-recursive globs, while GROMACS writes checkpoints and results under em/ and prod/. Consequently, clean()
leaves .cpt, .log, .xvg, and related outputs behind; a subsequent mdrun -deffnm can resume or reuse stale data. GROMACS cleanup must remove those
nested artifacts or directories.
- [P2] Convert statistical inefficiency using the output interval — a3fe/run/simulation.py:562
Gradient samples are spaced by dt * nstdhdl, but GradientData converts sample counts using only engine_config.timestep. Against current upstream
main, which includes the fs→ns conversion, this still underreports GROMACS inefficiency by nstdhdl—100× with defaults. On the unre-based feature tip,
the missing fs→ns conversion instead makes it 10,000× too large. Derive the interval from returned sample times or explicitly use dt * nstdhdl / 1000
ns.
Validate science:
Reviewed feature-gromacs-support against michellab/main at 62cf024. No files changed.
## Silently wrong
- a3fe/analyse/mbar.py:164: GROMACS MBAR rejects subsampling=True, then treats every correlated trajectory frame as independent. On the supplied
dataset, decorrelation reduced 501 frames/state to 19–27 and increased the two run uncertainties from 0.269/0.282 to 1.196/1.165 kcal/mol. Thus the
current errors are materially underestimated. Verified numerically; the point estimates’ scientific accuracy still requires domain validation.
- a3fe/run/stage.py:864: pymbar’s asymptotic 1σ errors are printed immediately before the claim that they are 95% confidence intervals. For supplied
run 1, 0.269 kcal/mol is reported where a normal 95% conversion would be approximately 0.527 kcal/mol—before correcting for autocorrelation. Verified
from pymbar semantics and the included data.
## Unvalidated
- a3fe/tests/test_analyse.py:136: the GROMACS regression values and underlying XVG dataset were introduced together, with no recorded independent
implementation, analytic result, or published reference. These tests would preserve an MBAR bug by updating the expected values. Repository
provenance verified; scientific correctness needs an independent oracle.
- a3fe/configuration/engine_config.py:641: only the discharge schedule is scientifically exercised. Restraint and vanish lambda mappings, endpoint
signs, and the combined ABFE cycle have no analytic or differential tests. The integration test merely requires finite values. Verified coverage gap;
stage conventions need domain confirmation.
- a3fe/configuration/engine_config.py:439: numerous physical/protocol defaults are described as matching fragment-opt-abfe-benchmark, but no pinned
source or comparison test establishes that provenance. Protocol drift would silently change results. Verified documentation gap; constants need a
domain expert to trace to the intended protocol.
## Fragile
- a3fe/analyse/mbar.py:192: analysis trusts the supplied/configured temperature without checking the temperature embedded in each XVG header. A stale
or modified Stage.pkl therefore produces plausible results using the wrong β. Verified behavior; actual-temperature policy needs confirmation.
- a3fe/analyse/mbar.py:195 and a3fe/analyse/mbar.py:596: MBAR states come from XVG column metadata, while reported lambda labels come independently
from numerically sorted directory names. Descending or non-monotonic custom schedules are accepted but can mislabel PMFs and neighboring gradients.
Verified by code-path inspection; a synthetic permutation invariant is missing.
Focused GROMACS analysis/read tests passed: 3 passed, 2 skipped. The full suite showed no failures during roughly three minutes but was manually
stopped before completion. There is no external ground truth here, so I cannot certify the physical results—only internal consistency and oracle
quality.
Before writing tests, the key invariants to confirm are: GROMACS errors should use decorrelated samples; per-run errors should remain explicitly 1σ or
be converted to 95%; and lambda schedules should be strictly increasing.
|
Also the |
|
OK, thinking about this a bit more, I think I've recommended moving too much stuff to A solution could be to create new |
|
Thank you Finlay for all your hard work! I will go through them carefully. |
2 participants
Description
Adds GROMACS engine support to a3fe, enabling non-adaptive ABFE calculations using GROMACS alongside the existing SOMD support. SOMD remains the default engine and continues to provide the adaptive workflow.
Key Changes
GromacsConfigfor MDP parameter management.Validation
88 passed, 10 skipped.Current Scope
GROMACS supports non-adaptive production calculations. Adaptive GROMACS execution is not implemented; the existing SOMD adaptive workflow is unchanged.