Add WiSE electrolyte benchmark (density, X-ray S(q), Li-O RDF) - #445
Add WiSE electrolyte benchmark (density, X-ray S(q), Li-O RDF)#445LucaBrugnoli wants to merge 13 commits into
Conversation
|
@joehart2001 The code and analysis are complete. The remaining step I think is uploading the trajectory data to |
Hi @LucaBrugnoli thanks for the PR! The easiest way is to attach a zip file containing your data to this PR and I can upload it. Hopefully that works |
|
Here are the 6 data files, one per model.
The expected directory structure on S3 is: Let me know if anything looks off or if you need the data in a different format mace-mh-1-omat.zip |
ElliottKasoar
added
the
new benchmark
dba4e00 to
ef5aecc
Compare
| return results | ||
|
|
||
|
|
||
| def normalize_metric(value: float, good: float, bad: float) -> float: |
There was a problem hiding this comment.
this will be taken care of when you build the table using the decorator. you could define your own function if you dont want to use our defualt, see the docs
| APP_ROOT = Path(__file__).resolve().parents[3] / "app" | ||
| OUT_PATH = APP_ROOT / "data" / "wise_electrolytes" / "density" | ||
|
|
||
| MODELS = [ |
There was a problem hiding this comment.
usually we import this to get all the models. this is better for in the future when we have more models than now.
from ml_peg.models.get_models import load_models
from ml_peg.models.models import current_models
| # --- Metrics table ----------------------------------------------------------- | ||
|
|
||
|
|
||
| def build_metrics_table(data: dict[str, dict]) -> dict: |
There was a problem hiding this comment.
we have decorators to build this automatically, see the tutorial
|
Hi @LucaBrugnoli thanks for the PR! From what i understand, you've provided us with NVT trajectories for each model and then your calc script takes these trajectories to calculate e.g. the rdf etc. Thank you for providing these trajectories! However, to make this test not rely on you computing these for each new model that is added in the future (as im sure you would prefer us running them for you!), I think we need to do some reshuffling. The ideal workflow would be:
analysis
app
|
- Merge density / rdf / xray_sf into a single litfsi_h2o_21m benchmark
under ml_peg/{calcs,analysis,app}/wise_electrolytes/litfsi_h2o_21m/.
- Adopt the standard ml-peg patterns in the analysis and calc scripts:
load_models(current_models) for model discovery, @build_table for
the metrics table, and metrics.yml for thresholds/tooltips/weights.
- Add a Janus recast of the LAMMPS+symmetrix Adastra production protocol
in ml_peg/calcs/wise_electrolytes/md_reference/calc_md_reference.py
(pytest-skipped reference; documents the exact MD parameters).
- Add the docs page docs/source/user_guide/benchmarks/wise_electrolytes.rst
(and toctree entry) and wire the app docs_url to it.
- Update the parent wise_electrolytes.yml to a single benchmark weight.
- Switch NPT (Melchionna) -> NPT_MTK (Martyna-Tobias-Klein) so the janus reference matches the LAMMPS fix npt formulation used in production. Pass thermostat_chain=3 and barostat_chain=3 explicitly, matching the LAMMPS default chain length. - Replace pytest.mark.skip with pytest.mark.very_slow per Joseph's review on PR ddmms#445: the reference protocol is now opt-in via --run-very-slow rather than unconditionally skipped, so it can be exercised when validating new models.
|
Thanks for the code updates! I think this test could also be well suited in the molecular dynamics category instead of its own one as its quite specific, opinions @ElliottKasoar? |
|
Thanks for both messages, I'll let you and @ElliottKasoar decide on the location; happy with whichever directory you prefer. |
- Merge density / rdf / xray_sf into a single litfsi_h2o_21m benchmark
under ml_peg/{calcs,analysis,app}/wise_electrolytes/litfsi_h2o_21m/.
- Adopt the standard ml-peg patterns in the analysis and calc scripts:
load_models(current_models) for model discovery, @build_table for
the metrics table, and metrics.yml for thresholds/tooltips/weights.
- Add a Janus recast of the LAMMPS+symmetrix Adastra production protocol
in ml_peg/calcs/wise_electrolytes/md_reference/calc_md_reference.py
(pytest-skipped reference; documents the exact MD parameters).
- Add the docs page docs/source/user_guide/benchmarks/wise_electrolytes.rst
(and toctree entry) and wire the app docs_url to it.
- Update the parent wise_electrolytes.yml to a single benchmark weight.
- Switch NPT (Melchionna) -> NPT_MTK (Martyna-Tobias-Klein) so the janus reference matches the LAMMPS fix npt formulation used in production. Pass thermostat_chain=3 and barostat_chain=3 explicitly, matching the LAMMPS default chain length. - Replace pytest.mark.skip with pytest.mark.very_slow per Joseph's review on PR ddmms#445: the reference protocol is now opt-in via --run-very-slow rather than unconditionally skipped, so it can be exercised when validating new models.
f3f66f7 to
b3816eb
Compare
Three sub-benchmarks for 21 m LiTFSI/H2O with 6 MLIP models (matpes-r2scan, mace-mpa-0-medium, mace-omat-0-medium, mace-mp-0b3, mace-mh-1-omat, mace-mh-1-omol): - NPT density vs Gilbert et al. JCED 2017 - X-ray structure factor S(q) vs SAXS experiment - Li-O RDF coordination numbers vs Watanabe et al. JPCB 2021
- Merge density / rdf / xray_sf into a single litfsi_h2o_21m benchmark
under ml_peg/{calcs,analysis,app}/wise_electrolytes/litfsi_h2o_21m/.
- Adopt the standard ml-peg patterns in the analysis and calc scripts:
load_models(current_models) for model discovery, @build_table for
the metrics table, and metrics.yml for thresholds/tooltips/weights.
- Add a Janus recast of the LAMMPS+symmetrix Adastra production protocol
in ml_peg/calcs/wise_electrolytes/md_reference/calc_md_reference.py
(pytest-skipped reference; documents the exact MD parameters).
- Add the docs page docs/source/user_guide/benchmarks/wise_electrolytes.rst
(and toctree entry) and wire the app docs_url to it.
- Update the parent wise_electrolytes.yml to a single benchmark weight.
- Switch NPT (Melchionna) -> NPT_MTK (Martyna-Tobias-Klein) so the janus reference matches the LAMMPS fix npt formulation used in production. Pass thermostat_chain=3 and barostat_chain=3 explicitly, matching the LAMMPS default chain length. - Replace pytest.mark.skip with pytest.mark.very_slow per Joseph's review on PR ddmms#445: the reference protocol is now opt-in via --run-very-slow rather than unconditionally skipped, so it can be exercised when validating new models.
- SAVGOL_WINDOW: 27 → 5 in calc_litfsi_h2o_21m.py (physical width 0.54 Å⁻¹ → 0.10 Å⁻¹), matching TRAVIS smoothing for dynasor/TRAVIS agreement - Add mace-mh-1-omat and mace-mh-1-omol to models.yml - Reset cubic cell to diagonal form after FIRE minimisation in md_reference
b3816eb to
e198469
Compare
There was a problem hiding this comment.
Can you revert these changes? We will test these models, but we make the changes locally (or would do this change as a separate PR, now that MACE has released these properly)
| model.default_dtype = "float64" | ||
| calc = model.get_calculator() |
There was a problem hiding this comment.
| model.default_dtype = "float64" | |
| calc = model.get_calculator() | |
| calc = model.get_calculator(precision="high") |
This is effectively the same, but more future-proof.
There was a problem hiding this comment.
Can this be combined with calc_litfsi_h2o_21m.py? Generally the idea is to have ml_peg/calcs/[category]/[benchmark]/calc_[benchmark].py,, but this isn't really its own benchmark?
We have similar examples e..g oxidation_states where there's a slow part that we may well run separately, and a faster part.
As long as the slow part is defined first, I think it's fine since they can be run either together or independently thanks to the very slow marker.
| if __name__ == "__main__": | ||
| import argparse | ||
|
|
||
| parser = argparse.ArgumentParser( | ||
| description="Run the WiSE reference MD protocol for one registered model." | ||
| ) | ||
| parser.add_argument("model", help=f"one of: {sorted(MODELS)}") | ||
| parser.add_argument( | ||
| "--continuation", | ||
| action="store_true", | ||
| help="append the optional 50 ps NVT continuation after NPT.", | ||
| ) | ||
| args = parser.parse_args() | ||
|
|
||
| if args.model not in MODELS: | ||
| parser.error(f"unknown model '{args.model}'. Registered: {sorted(MODELS)}") | ||
|
|
||
| run_reference_md(args.model, MODELS[args.model], run_continuation=args.continuation) |
There was a problem hiding this comment.
A couple of questions/suggestions:
- What's the context of
continuation- when would we/wouldn't we want to use it? Generally, I think it's good to define a single protocol for all the models - It's preferable to integrate with the
ml_peg calcCLI, which also allows us to add custom inputs if you wanted to add an new conftest.py (see the water and liquid densities in the MD category for examples). That already has built in things for choosing the model(s) etc.
There was a problem hiding this comment.
Thanks you both for the detailed review, and apologies for the delayed response. Here's a summary of what I've changed, plus two open questions I can't settle on my own.
-
models.yml
Reverted to the upstream state, as requested. This undoes the model additions that went in with e198469; ok for the MH-1 heads to be enabled on your side or in a separate PR now that MACE has released them properly.
My results for those two models don't depend on the registry flag, so nothing is lost by leaving them commented. -
Reference MD protocol
Your question about continuation prompted me to re-check the script against the protocol that actually produced the published trajectories, and I found two discrepancies. The script was running minimization → 50 ps NVT → 200 ps NPT on p64_w170, with an optional NVT continuation, and it never used p16_w42. This differs from the production workflow in two important ways: it included an NPT stage that was not used in production, and, because the density data were obtained from the smaller p16_w42 cell, it could not have regenerated density.json.
This is what production actually did, and what the script now does:
- p64_w170 (S(q) and RDF): Min → NVT 50 ps equilibration → NVT 50 ps production, held at the experimental volume (L =27.49 Å) throughout. No NPT: the box is fixed at the experimental density, so that S(q) and the coordination numbers are compared at the reference density rather than at each model's own.
- p16_w42 (density): Min → NVT 50 ps → NPT 200 ps, with the density averaged over the last 150 ps. NPT is used only here, on the smaller cell, because it's cheaper to converge.
One consequence worth flagging: the mace-mp-0b3 density I sent in April came from a 50 ps NPT run, so it no longer matched the protocol above. I've rerun it at 200 ps like the other three. The density was practically unchanged (from 1.5632 to 1.5623 g/cm³), but the set is now homogeneous. The updated zip is attached below.
-
Continuation
This option was a consequence of the protocol mismatch described above. The second 50 ps NVT block is not an optional continuation after NPT, but the production stage that follows equilibration directly, with no intervening NPT stage. Once the protocol was corrected, the flag no longer served a purpose, so I removed it. All models now follow the same single protocol.
When a test is rerun, it instead detects and skips any stages that have already reached their final step, allowing interrupted runs to resume without repeating the entire workflow. -
Scripts combined
I merged the two scripts into a single calc_litfsi_h2o_21m.py, following the structure used in oxidation_states.
The very_slow MD test is defined first and generates the trajectories. The extraction tests then read those trajectories and, when the required data are unavailable, skip with a reference to the MD test that produces them. The MD and analysis stages can therefore be run either together or independently. The separate md_reference/ directory has been removed, leaving the benchmark contained in a single directory as intended. -
Precision
Applied as suggested. I also removed the model.default_dtype = "float64" assignment immediately above it.
About the earlier review comments
@joehart2001, the three comments from April are now addressed in the consolidated analyse_litfsi_h2o_21m.py. It uses the build_table decorator for table generation and scoring, and loads the registry through load_models(current_models) rather than using a hard-coded model list.
analyse_density.py has been removed, so I believe those three threads can now be resolved.
Category change
I agree that the electrolytes category introduced in PR (#394) is the most appropriate location for this benchmark. Since that PR has not yet been completed, would you prefer me to restructure the benchmark against it now, or leave the current location unchanged and move it once PR (#394) is merged?
Additional data upload
The April upload contained only the generated outputs, not the initial structures. I have attached the missing inputs here, together with the updated mace-mp-0b3 results:
litfsi_h2o_21m_initial_structures.zip: contains p64_w170_initial.xyz (1,534 atoms) and p16_w42_initial.xyz (382 atoms), both at the experimental density.
mace-mp-0b3.zip: replaces the archive uploaded in April. It contains the same NVT trajectory, together with the updated density.json from the 200 ps rerun.
The MH-1 heads were enabled here to run them through this benchmark, but models.yml is shared by every benchmark in the repo, so activating them is not this PR's call to make.
The reference MD did not match the protocol that produced the published
data. It ran NPT 200 ps on p64_w170 and never ran p16_w42 at all, so it
sampled S(q) and the RDF at each model's own density rather than at the
experimental one, and could not have regenerated the density data.
The production protocol, and what this now does:
p64_w170 Min -> NVT 50 ps eq -> NVT 50 ps prod at the experimental
volume. No NPT. Gives nvt_trajectory.extxyz.
p16_w42 Min -> NVT 50 ps -> NPT 200 ps, density averaged over the
last 150 ps. Gives density.json.
The continuation flag was an artefact of the same mis-transcription: the
second NVT block is the production stage, not an optional addition after
NPT. With the protocol corrected there is nothing left for it to select,
so it is gone and the protocol is identical for every model. Re-running
the test now skips whichever stages already reached their final step.
Minimization passes filter_class=None so that only the atoms relax:
janus-core defaults to FrechetCellFilter, which would have relaxed the
cell away from the experimental volume the protocol depends on.
Both halves live in calc_litfsi_h2o_21m.py, as in calc_oxidation_states.py:
the very_slow MD test is defined first and writes the data, the extraction
tests read it back and skip with a pointer to the MD test if it is absent.
Also switch to get_calculator(precision="high") and drop the preceding
model.default_dtype assignment, which took precedence over the argument
and would have made the change a no-op.
current_models moved to ml_peg/models/__init__.py in ddmms#560, so importing it from ml_peg.models.models raised ImportError and every test in this benchmark errored at collection. All other calcs use `from ml_peg.models import current_models`; do the same.
The Savitzky-Golay window is 5, not 27: at dq = 0.02 A^-1 that is the 0.10 A^-1 physical width the calc applies. The 27 referred to an earlier, finer q grid. Also describe the protocol per cell, and point at the reference MD in its new home now that md_reference/ is gone.
figure_density_timeseries was loaded but never mapped to a table column, so it was never displayed, and ruff rejects the dead assignment.
CI runs the full pre-commit suite, which this benchmark had never been through: ruff-format wanted to reformat both files, and numpydoc rejected two parameter descriptions that opened with a lowercase "janus-core". The reformatting is cosmetic only; the protocol constants, Cromer-Mann form factors and composition are unchanged.
|
Reposting here because my earlier reply ended up in a review thread on Thank you both for the detailed review, and apologies for the delayed response. Here's a summary of what I've changed, plus two open questions I can't settle on my own.
This is what production actually did, and what the script now does:
One consequence worth flagging: the mace-mp-0b3 density I sent in April came from a 50 ps NPT run, so it did not match the protocol above. I've rerun it for 200 ps, as for the other three models. The density was practically unchanged (from 1.5632 to 1.5623 g/cm³), but the dataset is now consistent. The updated zip is attached below.
About the earlier review comments Category change Additional data upload The April upload contained only the generated outputs, not the initial structures, so the MD test cannot currently be run without the local input files. I have attached the missing inputs here, together with the updated mace-mp-0b3 results: litfsi_h2o_21m_initial_structures.zip: contains p64_w170_initial.xyz (1,534 atoms) and p16_w42_initial.xyz (382 atoms), both at the experimental density. litfsi_h2o_21m_initial_structures.zip The calculation still reads its inputs from a local directory rather than through download_s3_data, which is the last remaining change I am aware of. What would be the best way to handle this from here? The pre-commit failure shown for 6a972ba was fixed in 4cd1ff2, including the ruff-format changes and two NumPy-style docstring fixes. The CI run for the latest commit is currently waiting for approval. |
Relocate calcs, analysis, app and docs from wise_electrolytes/ to electrolytes/, following the category introduced in ddmms#394.
Replace the experimental per-species CN references (2.0/2.0, Watanabe et al. 2021, measured at ~18.5 m) with theory-anchored values at the benchmark composition: 2.4 (water), 1.9 (TFSI) and 4.3 (total), from r2SCAN AIMD and two independent MACE potentials trained on the same r2SCAN dataset (converged over >2 ns). The experimental total (4.21) is kept as a consistency check; the experimental water/TFSI split is dropped, as it is at a different concentration and disagrees systematically with all MD-based partitionings. Add a total-CN metric and rebalance weights (0.25/0.25/0.5) so the CN family carries the same total weight as density and S(q); good thresholds now equal the reference uncertainties (0.15/0.05/0.10).
Models trained on molecular data refuse to evaluate a structure whose info dict lacks these keys: orb-v3-consv-omol raises "atoms.info must contain both 'charge' and 'spin'", and the same applies to the other OMol-headed models in models.yml. Both cells are neutral and closed shell, so the defaults used by the other calculation scripts (charge 0, spin 1) are the right ones here too.
3 participants
Pre-review checklist for PR author
Summary
New benchmark for 21 molal LiTFSI/H₂O water-in-salt electrolyte (WiSE), evaluating MLIP foundation models on three
experimental observables:
(2017)
(2021). Computed via dynasor.
(2021)
Further details on the simulation protocol and MLIP assessment for this system: L. Brugnoli, arXiv:2603.22099
(2026).
Linked issue
Resolves #304
Progress
Note: Trajectory data (~500 MB, extxyz) needs to be uploaded to the ml-peg S3 bucket. Data files are available on
request.
Testing
Tested on 6 models:
matpes-r2scan,mace-mpa-0-medium,mace-omat-0-medium,mace-mp-0b3,mace-mh-1-omat,mace-mh-1-omol.Requirement: ASE < 3.28 (3.28.0 has a bug in
ase.io.extxyz.ixyzchunks). Tests run with--noconftest.New decorators/callbacks
No new callbacks required. The RDF app uses existing
plot_from_table_columnfrom ml-peg utils.