diff --git a/docs/source/user_guide/benchmarks/electrolytes.rst b/docs/source/user_guide/benchmarks/electrolytes.rst new file mode 100644 index 000000000..6c194cb21 --- /dev/null +++ b/docs/source/user_guide/benchmarks/electrolytes.rst @@ -0,0 +1,118 @@ +============ +Electrolytes +============ + +LiTFSI/H2O 21 m +=============== + +Summary +------- + +Performance in predicting structural and thermodynamic properties of the 21 molal +LiTFSI/H2O "water-in-salt" electrolyte (WiSE). The benchmark consolidates three +experimental observables for the same chemistry: NPT equilibrium density, Li-O +coordination from the radial distribution function, and the X-ray static +structure factor S(q). + +Two simulation cells are used: a small cell (16 LiTFSI + 42 H2O, 382 atoms) for +the NPT density, and a large cubic cell (64 LiTFSI + 170 H2O, 1534 atoms, +27.4938 A) for the NVT trajectories that feed the RDF and S(q). + +Metrics +------- + +1. Density error + + Absolute error in NPT density vs the experimental value + (1.7126 g/cm3, Gilbert et al., *J. Chem. Eng. Data* 62, 2056 (2017)). + The density is the average over the production NPT trajectory at + 298.15 K and 1 atm. + +2. Li-O_water coordination number error + + Absolute error in the Li-O_water coordination number, computed by + integrating the Li-O_water radial distribution function up to the first + minimum (R_cut = 2.83 A, from r2SCAN AIMD reference). Water oxygens are + identified from the first frame as those bonded to a hydrogen + (d_OH < 1.25 A). Reference value: 2.4, anchored to theory at 21 m + (r2SCAN AIMD, Li128, 20 ps: 2.23; MACE potentials trained on the same + r2SCAN dataset, converged over >2 ns: 2.35-2.39). The experimental + partitioning of Watanabe et al., *J. Phys. Chem. B* 125, 7477 (2021) + (1.93 water / 2.28 TFSI, neutron diffraction with isotopic substitution) + is measured at a different composition (~18.5 m) and disagrees + systematically with all MD-based partitionings (classical, AIMD, and + MLIP), suggesting a method-related bias; it is therefore not used for + the per-species references. Weight: 0.25; the "good" threshold (0.15) + equals the uncertainty of the reference (AIMD vs converged-MLIP spread). + +3. Li-O_TFSI coordination number error + + Same definition as above, restricted to TFSI oxygens (those bonded to a + sulfur, d_OS < 1.75 A). Reference value: 1.9, anchored to theory at 21 m + (r2SCAN AIMD: 2.02; converged r2SCAN-trained MACE: 1.89-1.91). + Weight: 0.25; "good" threshold 0.05 (FT/TfS agreement). + +4. Li-O_total coordination number error + + Absolute error in the total Li-O coordination number (water + TFSI). + Reference value: 4.3, from theory at 21 m (r2SCAN AIMD: 4.31; converged + r2SCAN-trained MACE: 4.27-4.29), and consistent with the experimental + total of 4.21 +/- 0.03 at ~18.5 m (Watanabe et al. 2021), which is + robust to the concentration difference. Weight: 0.5; "good" threshold + 0.10 (theory-experiment spread). The three CN metrics together carry + the same total weight (1.0) as the density and S(q) families. + +5. S(q) R-factor + + R-factor between the computed and experimental X-ray structure factor, + ``sum(|S_exp - S_calc|) / sum(|S_exp|)``. S(q) is computed via dynasor + in the Faber-Ziman convention with Cromer-Mann 4-Gaussian form factors + (including hydrogen) and Savitzky-Golay smoothing + (window = 5, order = 3, dq = 0.02 A^-1, i.e. a physical width of + 0.10 A^-1). Reference: SAXS data of + Zhang et al., *J. Phys. Chem. B* 125, 4501 (2021). + +Computational cost +------------------ + +High: production trajectories require a GPU MD code (e.g. LAMMPS + symmetrix +or Janus + MACE) and several hours per model on a single MI250X / A100 GPU. +Re-extracting metrics from pre-computed trajectories is cheap (seconds per +model). + +Data availability +----------------- + +Input structures: + +* p64_w170 cubic cell (1534 atoms, 27.4938 A) for NVT trajectories. +* p16_w42 cell (382 atoms) for NPT density. + +The reference data were generated with LAMMPS + symmetrix on Adastra (MI250X) +using the protocol below, at dt = 0.5 fs, T = 298.15 K, P = 1.01325 bar, +Nose-Hoover thermostat tau = 50 fs and barostat tau = 500 fs: + +* p64_w170: Min -> NVT 50 ps equilibration -> NVT 50 ps production, held at the + experimental volume throughout (no NPT), giving the trajectory used for S(q) + and the RDF. +* p16_w42: Min -> NVT 50 ps -> NPT 200 ps, with the density averaged over the + last 150 ps. + +The janus-core recast of the same protocol is ``test_reference_md`` in +``ml_peg/calcs/electrolytes/litfsi_h2o_21m/``, marked ``very_slow``. It +regenerates both data products for any registered model, so the benchmark does +not depend on the pre-computed trajectories. + +Reference data: + +* Density: Gilbert et al., *J. Chem. Eng. Data* 62, 2056 (2017), + DOI: 10.1021/acs.jced.7b00135. +* Li-O coordination numbers: theoretical references at 21 m (r2SCAN AIMD and + r2SCAN-trained MACE potentials, Brugnoli et al., arXiv:2603.22099); + experimental total CN at ~18.5 m: Watanabe et al., *J. Phys. Chem. B* 125, + 7477 (2021), DOI: 10.1021/acs.jpcb.1c04693. +* X-ray S(q): Zhang et al., *J. Phys. Chem. B* 125, 4501 (2021), + DOI: 10.1021/acs.jpcb.1c02189. + +Further details on the simulation protocol and MLIP assessment for this +system: Brugnoli et al., arXiv:2603.22099 (2026). diff --git a/docs/source/user_guide/benchmarks/index.rst b/docs/source/user_guide/benchmarks/index.rst index a5a4e6cd1..4b8506628 100644 --- a/docs/source/user_guide/benchmarks/index.rst +++ b/docs/source/user_guide/benchmarks/index.rst @@ -14,6 +14,7 @@ Benchmarks bulk_crystal lanthanides non_covalent_interactions + electrolytes tm_complexes conformers molecular_dynamics diff --git a/ml_peg/analysis/electrolytes/data/saxs_maginn.txt b/ml_peg/analysis/electrolytes/data/saxs_maginn.txt new file mode 100644 index 000000000..08d66b2cf --- /dev/null +++ b/ml_peg/analysis/electrolytes/data/saxs_maginn.txt @@ -0,0 +1,301 @@ +# 10.1021/acs.jpcb.1c02189 +# q / A-1 S(q) +0.20849 -0.14530 +0.27792 -0.13506 +0.34735 -0.12067 +0.41678 -0.10336 +0.45728 -0.08727 +0.49200 -0.06776 +0.52671 -0.04801 +0.55178 -0.02997 +0.57493 -0.00949 +0.59036 0.01014 +0.61598 0.03412 +0.64243 0.06049 +0.65400 0.08160 +0.66557 0.10150 +0.67714 0.12052 +0.68293 0.13778 +0.69797 0.15855 +0.70800 0.17996 +0.72012 0.20988 +0.73170 0.24081 +0.74465 0.27042 +0.75484 0.30330 +0.73500 0.28700 +0.76641 0.34761 +0.76586 0.32650 +0.77798 0.38731 +0.77550 0.36819 +0.79286 0.42909 +0.78708 0.40842 +0.80113 0.46380 +0.78129 0.44792 +0.81022 0.50741 +0.81215 0.48840 +0.82427 0.56746 +0.82526 0.54096 +0.81601 0.52692 +0.83584 0.62473 +0.83336 0.60080 +0.81601 0.59129 +0.84741 0.67948 +0.84493 0.65675 +0.82758 0.64688 +0.85899 0.73424 +0.85651 0.71161 +0.83915 0.70247 +0.87001 0.79805 +0.86808 0.77562 +0.86808 0.75989 +0.87965 0.86242 +0.88312 0.83911 +0.87386 0.82146 +0.89370 0.93738 +0.89122 0.91350 +0.88544 0.89753 +0.88544 0.88241 +0.90527 0.99757 +0.89701 0.97555 +0.89701 0.96141 +0.91436 1.05943 +0.90858 1.03504 +0.90858 1.01993 +0.92842 1.12024 +0.92015 1.09454 +0.92015 1.07991 +0.93999 1.17479 +0.93172 1.15208 +0.92015 1.14135 +0.95156 1.21659 +0.94329 1.19840 +0.96258 1.25522 +0.95487 1.23790 +0.97470 1.28555 +0.98727 1.31134 +1.01272 1.33234 +1.05323 1.30934 +1.04744 1.33153 +1.06595 1.28501 +1.07637 1.26497 +1.08794 1.24156 +1.10067 1.21596 +1.11108 1.19365 +1.12150 1.17119 +1.13307 1.14808 +1.14580 1.12380 +1.15853 1.10039 +1.17010 1.07815 +1.18630 1.05913 +1.19980 1.03724 +1.23259 1.01286 +1.27226 0.99067 +1.34252 0.98092 +1.41195 0.97409 +1.46237 0.95577 +1.49874 0.93595 +1.51417 0.92045 +1.53577 0.90319 +1.54965 0.88446 +1.57395 0.86291 +1.59710 0.83755 +1.62024 0.81219 +1.64338 0.78586 +1.66653 0.76026 +1.69380 0.73737 +1.71667 0.71637 +1.75332 0.69589 +1.79299 0.67865 +1.86325 0.66980 +1.93846 0.67614 +1.99748 0.69413 +2.04261 0.71832 +2.07154 0.73685 +2.09468 0.75441 +2.12554 0.77245 +2.14868 0.79293 +2.17183 0.81439 +2.19497 0.83584 +2.21040 0.85462 +2.23354 0.87558 +2.26247 0.90216 +2.27983 0.92593 +2.30297 0.94776 +2.33190 0.97567 +2.35697 1.00091 +2.37587 1.01876 +2.39555 1.03748 +2.42448 1.05796 +2.46498 1.08027 +2.52284 1.09893 +2.59227 1.09722 +2.64665 1.08137 +2.68484 1.06363 +2.72534 1.04480 +2.76502 1.02682 +2.79477 1.01056 +2.82949 0.99506 +2.86420 0.97750 +2.89892 0.95995 +2.93859 0.94156 +2.96835 0.92630 +3.00885 0.90893 +3.05513 0.88918 +3.10060 0.87154 +3.15019 0.85357 +3.19400 0.83804 +3.25764 0.82126 +3.29236 0.80890 +3.34443 0.79342 +3.40807 0.77854 +3.47172 0.76148 +3.54115 0.74636 +3.61058 0.73295 +3.68001 0.72198 +3.74944 0.71491 +3.81887 0.71076 +3.88830 0.71101 +3.95773 0.71539 +4.02716 0.72612 +4.09080 0.74051 +4.14866 0.75709 +4.19660 0.77541 +4.23958 0.79547 +4.27132 0.81512 +4.30901 0.83455 +4.34075 0.85491 +4.37431 0.87607 +4.40787 0.89792 +4.43217 0.91752 +4.46303 0.93459 +4.48617 0.95214 +4.50932 0.97019 +4.53246 0.98921 +4.55560 1.00822 +4.57875 1.02724 +4.60189 1.04724 +4.62503 1.06698 +4.64818 1.08747 +4.66708 1.10595 +4.69446 1.12623 +4.71761 1.14476 +4.74075 1.16305 +4.76389 1.18231 +4.78511 1.20279 +4.80825 1.22298 +4.83718 1.24193 +4.86611 1.26180 +4.90083 1.28374 +4.93554 1.30447 +4.97604 1.32440 +5.02150 1.34407 +5.08019 1.36420 +5.14962 1.37444 +5.21905 1.37225 +5.28269 1.36225 +5.34055 1.34558 +5.38105 1.32586 +5.42155 1.30788 +5.45627 1.29057 +5.49098 1.27301 +5.52570 1.25546 +5.56041 1.23717 +5.59513 1.21742 +5.62985 1.19694 +5.66456 1.17719 +5.69928 1.15671 +5.73399 1.13647 +5.76871 1.11721 +5.80342 1.09795 +5.83814 1.07991 +5.87285 1.06235 +5.90757 1.04650 +5.94724 1.02975 +5.99353 1.01199 +6.06379 0.99359 +6.13322 0.97799 +6.20265 0.96604 +6.27208 0.96043 +6.34151 0.95726 +6.41094 0.95775 +6.48037 0.96043 +6.54980 0.96507 +6.61923 0.97043 +6.68866 0.97458 +6.75809 0.97897 +6.82752 0.98140 +6.89695 0.98287 +6.96638 0.98213 +7.03581 0.98067 +7.10524 0.97823 +7.17467 0.97531 +7.24410 0.97238 +7.31353 0.96872 +7.38296 0.96580 +7.45239 0.96287 +7.52182 0.96043 +7.59125 0.95775 +7.66068 0.95629 +7.73011 0.95434 +7.79954 0.95288 +7.86898 0.95214 +7.93841 0.95117 +8.00784 0.95214 +8.07727 0.95263 +8.14670 0.95385 +8.21613 0.95653 +8.28556 0.96019 +8.35499 0.96458 +8.42442 0.97019 +8.49385 0.97604 +8.56328 0.98238 +8.63271 0.98896 +8.70214 0.99530 +8.77157 1.00262 +8.84100 1.00944 +8.91043 1.01603 +8.97986 1.02432 +9.04929 1.03261 +9.11872 1.04163 +9.18815 1.05162 +9.25758 1.06284 +9.32701 1.07527 +9.39644 1.08747 +9.46587 1.10088 +9.53530 1.11404 +9.60473 1.12745 +9.67417 1.14038 +9.74360 1.15305 +9.81303 1.16329 +9.88246 1.17158 +9.95189 1.17646 +10.02132 1.18061 +10.09075 1.17987 +10.16018 1.17670 +10.22961 1.16890 +10.29904 1.15744 +10.36847 1.14184 +10.43211 1.12584 +10.48997 1.10917 +10.54204 1.08887 +10.60238 1.06946 +10.65583 1.05065 +10.70488 1.03288 +10.76046 1.01389 +10.81629 0.99652 +10.87184 0.98043 +10.92970 0.96346 +10.99334 0.94776 +11.06277 0.93240 +11.13220 0.92069 +11.20163 0.91265 +11.27106 0.90557 +11.34049 0.90314 +11.40992 0.90289 +11.47936 0.90411 +11.54879 0.90704 +11.61822 0.91167 +11.68765 0.91703 +11.75708 0.92362 +11.82651 0.92996 +11.89594 0.93678 diff --git a/ml_peg/analysis/electrolytes/litfsi_h2o_21m/analyse_litfsi_h2o_21m.py b/ml_peg/analysis/electrolytes/litfsi_h2o_21m/analyse_litfsi_h2o_21m.py new file mode 100644 index 000000000..ef387bb87 --- /dev/null +++ b/ml_peg/analysis/electrolytes/litfsi_h2o_21m/analyse_litfsi_h2o_21m.py @@ -0,0 +1,733 @@ +"""Analyse the WiSE 21 m LiTFSI/H2O electrolyte benchmark.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import plotly.graph_objects as go +import pytest +from scipy.interpolate import interp1d + +from ml_peg.analysis.utils.decorators import build_table +from ml_peg.analysis.utils.utils import load_metrics_config +from ml_peg.app import APP_ROOT +from ml_peg.calcs import CALCS_ROOT +from ml_peg.models import current_models +from ml_peg.models.get_models import load_models + +# --- Paths ------------------------------------------------------------------- + +CALC_PATH = CALCS_ROOT / "electrolytes" / "litfsi_h2o_21m" / "outputs" +OUT_PATH = APP_ROOT / "data" / "electrolytes" / "litfsi_h2o_21m" +EXP_SAXS_PATH = Path(__file__).resolve().parents[1] / "data" / "saxs_maginn.txt" + +MODELS = load_models(current_models) + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + +# --- Physical constants and reference values --------------------------------- + +# Box geometry (same for all models: p64_w170 NVT cell) +L_BOX = 27.4938 # Å (cubic NVT cell at experimental density) +V_BOX = L_BOX**3 # ų + +# Integration cutoff: first minimum of Li-O_total g(r) (r2SCAN AIMD) +R_CUT = 2.83 # Å +DR = 0.02 # Å (bin width used in calc_rdf.py) + +# Experimental references +RHO_EXP = 1.7126 # g/cm³ — Gilbert et al., JCED 62, 2056 (2017) + +# Coordination number references. +# The total CN is anchored to theory at 21 m: r2SCAN AIMD (Li128, 20 ps) gives +# 4.31, and two independent MACE potentials trained on the same r2SCAN dataset +# (fine-tuned and from-scratch) converge to 4.27-4.29 over >2 ns; the neutron +# diffraction experiment (Watanabe et al., JPCB 125, 7477, 2021) gives 4.21 at +# the nearby composition x_LiTFSA=0.25 (~18.5 m). The water/TFSI partitioning +# is anchored to theory only: the experimental split (1.93/2.28) is at a +# different concentration and shows a systematic inversion vs all MD methods, +# so it is not used. AIMD (20 ps): 2.23/2.02; converged r2SCAN-trained MACE +# (>2 ns): 2.35-2.39 (water) and 1.89-1.91 (TFSI). +CN_REF_WATER = 2.4 # r2SCAN AIMD + converged r2SCAN-trained MACE potentials +CN_REF_TFSI = 1.9 # r2SCAN AIMD + converged r2SCAN-trained MACE potentials +CN_REF_TOTAL = 4.3 # theory at 21 m; exp. 4.21 at ~18.5 m (Watanabe 2021) + + +# ============================================================================= +# Density helpers and figures +# ============================================================================= + + +def load_density_results() -> dict[str, dict]: + """ + Load density.json for each registered model that has calc output. + + Returns + ------- + dict[str, dict] + Per-model density results keyed by registry model name. + """ + results = {} + for model in MODELS: + json_path = CALC_PATH / model / "density.json" + if json_path.exists(): + with open(json_path) as f: + results[model] = json.load(f) + return results + + +def build_density_bar_chart(data: dict[str, dict]) -> go.Figure: + """ + Build bar chart of density per model. + + Parameters + ---------- + data : dict[str, dict] + Per-model density results. + + Returns + ------- + go.Figure + Plotly bar chart figure. + """ + fig = go.Figure() + for model in MODELS: + if model not in data: + continue + fig.add_trace( + go.Bar( + x=[model], + y=[data[model]["rho_mean"]], + error_y={"type": "data", "array": [data[model]["rho_std"]]}, + name=model, + ) + ) + + fig.add_hline( + y=RHO_EXP, + line_dash="dash", + line_color="red", + annotation_text=f"Exp. ({RHO_EXP} g/cm\u00b3, Gilbert 2017)", + ) + fig.update_layout( + title="Electrolyte Density (LiTFSI/H2O, 21 m)", + yaxis_title="Density / g cm⁻³", + showlegend=False, + ) + return fig + + +def build_density_timeseries(data: dict[str, dict]) -> go.Figure: + """ + Build density vs time plot for all models. + + Parameters + ---------- + data : dict[str, dict] + Per-model density results. + + Returns + ------- + go.Figure + Plotly timeseries figure. + """ + fig = go.Figure() + for model in MODELS: + if model not in data: + continue + d = data[model] + fig.add_trace( + go.Scatter( + x=d["time_full"], + y=d["density_full"], + mode="lines", + name=model, + opacity=0.8, + ) + ) + + fig.add_hline( + y=RHO_EXP, + line_dash="dash", + line_color="black", + annotation_text=f"Exp. ({RHO_EXP} g/cm\u00b3)", + ) + fig.update_layout( + title="NPT Density vs Time (LiTFSI/H2O)", + xaxis_title="Time / ps", + yaxis_title="Density / g cm⁻³", + ) + return fig + + +# ============================================================================= +# RDF / coordination number helpers and figures +# ============================================================================= + + +def compute_cn_from_gr( + gr: np.ndarray, + r: np.ndarray, + n_neighbor: int, + volume: float, + r_cut: float, + dr: float, +) -> float: + """ + Integrate g(r) to coordination number: CN = 4pi rho int_0^r_cut g(r) r^2 dr. + + Parameters + ---------- + gr : np.ndarray + Radial distribution function values. + r : np.ndarray + Radial distance values in angstrom. + n_neighbor : int + Number of neighbor atoms of the relevant species in the simulation box. + volume : float + Volume of the simulation box in angstrom^3. + r_cut : float + Integration cutoff distance in angstrom. + dr : float + Bin width in angstrom. + + Returns + ------- + float + Coordination number obtained by integrating g(r) up to ``r_cut``. + """ + rho = n_neighbor / volume + mask = r <= r_cut + return float(np.sum(4 * np.pi * rho * gr[mask] * r[mask] ** 2 * dr)) + + +def load_rdf_results() -> dict[str, dict]: + """ + Load rdf.json and rdf.npz for each registered model that has calc output. + + Returns + ------- + dict[str, dict] + Mapping from registry model name to a dict containing coordination + numbers, errors, radial distances, g(r) arrays, and frame count. + """ + results = {} + for model in MODELS: + json_path = CALC_PATH / model / "rdf.json" + npz_path = CALC_PATH / model / "rdf.npz" + if not json_path.exists() or not npz_path.exists(): + continue + + with open(json_path) as f: + meta = json.load(f) + + data = np.load(npz_path) + r = data["r"] + + n_ow = meta["n_O_water"] + n_of = meta["n_O_TFSI"] + + cn_w = compute_cn_from_gr(data["gr_LiO_water"], r, n_ow, V_BOX, R_CUT, DR) + cn_f = compute_cn_from_gr(data["gr_LiO_TFSI"], r, n_of, V_BOX, R_CUT, DR) + + results[model] = { + "cn_water": cn_w, + "cn_tfsi": cn_f, + "cn_total": cn_w + cn_f, + "err_water": abs(cn_w - CN_REF_WATER), + "err_tfsi": abs(cn_f - CN_REF_TFSI), + "err_total": abs(cn_w + cn_f - CN_REF_TOTAL), + "r": r, + "gr_water": data["gr_LiO_water"], + "gr_tfsi": data["gr_LiO_TFSI"], + "gr_total": data["gr_LiO_total"], + "n_frames": meta["n_frames_used"], + } + return results + + +def build_cn_bar_chart(data: dict[str, dict]) -> go.Figure: + """ + Bar chart of CN_water and CN_TFSI per model with experimental reference. + + Parameters + ---------- + data : dict[str, dict] + Per-model RDF results as returned by :func:`load_rdf_results`. + + Returns + ------- + go.Figure + Plotly figure with grouped bars and experimental reference lines. + """ + fig = go.Figure() + + x_models = [m for m in MODELS if m in data] + cn_water_vals = [data[m]["cn_water"] for m in x_models] + cn_tfsi_vals = [data[m]["cn_tfsi"] for m in x_models] + + fig.add_trace( + go.Bar( + name="Li-Owater", + x=x_models, + y=cn_water_vals, + marker_color="steelblue", + ) + ) + fig.add_trace( + go.Bar( + name="Li-OTFSI", + x=x_models, + y=cn_tfsi_vals, + marker_color="coral", + ) + ) + + fig.add_hline( + y=CN_REF_WATER, + line_dash="dash", + line_color="steelblue", + annotation_text=f"Ref. Owater ({CN_REF_WATER}, r2SCAN)", + annotation_position="top right", + ) + fig.add_hline( + y=CN_REF_TFSI, + line_dash="dash", + line_color="coral", + annotation_text=f"Ref. OTFSI ({CN_REF_TFSI}, r2SCAN)", + annotation_position="bottom right", + ) + + fig.update_layout( + title="Li⁺ Coordination Numbers (LiTFSI/H₂O, 21 m)", + yaxis_title="Coordination Number", + barmode="group", + legend={"x": 0.01, "y": 0.99}, + ) + return fig + + +def build_gr_plot(data: dict[str, dict]) -> go.Figure: + """ + G(r) plot: Li-O_water and Li-O_TFSI for all models. + + Parameters + ---------- + data : dict[str, dict] + Per-model RDF results as returned by :func:`load_rdf_results`. + + Returns + ------- + go.Figure + Plotly figure with g(r) curves and integration cutoff line. + """ + fig = go.Figure() + + colors = [ + "#1f77b4", + "#ff7f0e", + "#2ca02c", + "#d62728", + "#9467bd", + "#8c564b", + "#e377c2", + "#7f7f7f", + ] + + for i, model in enumerate(MODELS): + if model not in data: + continue + d = data[model] + c = colors[i % len(colors)] + fig.add_trace( + go.Scatter( + x=d["r"], + y=d["gr_water"], + mode="lines", + name=f"{model} Ow", + line={"color": c, "dash": "solid"}, + legendgroup=model, + ) + ) + fig.add_trace( + go.Scatter( + x=d["r"], + y=d["gr_tfsi"], + mode="lines", + name=f"{model} OTFSI", + line={"color": c, "dash": "dot"}, + legendgroup=model, + ) + ) + + fig.add_vline( + x=R_CUT, + line_dash="dash", + line_color="gray", + annotation_text=f"r_cut={R_CUT} Å", + annotation_position="top right", + ) + + fig.update_layout( + title="Li-O Radial Distribution Functions (LiTFSI/H₂O, 21 m)", + xaxis_title="r / Å", + yaxis_title="g(r)", + xaxis_range=[1.0, 6.0], + ) + return fig + + +# ============================================================================= +# X-ray S(q) helpers and figures +# ============================================================================= + + +def load_experimental_sq() -> tuple[np.ndarray, np.ndarray]: + """ + Load experimental SAXS S(q) data. + + Returns + ------- + q_exp : np.ndarray + Scattering vector values in inverse angstroms. + sq_exp : np.ndarray + Experimental structure factor values. + """ + data = np.loadtxt(EXP_SAXS_PATH) + return data[:, 0], data[:, 1] + + +def load_computed_sq(model: str) -> tuple[np.ndarray, np.ndarray] | None: + """ + Load computed S(q) for a model. + + Parameters + ---------- + model : str + Registry name of the MLIP model. + + Returns + ------- + tuple of (np.ndarray, np.ndarray) or None + Scattering vector and structure factor arrays, or None if the + output file does not exist. + """ + json_path = CALC_PATH / model / "xray_sq.json" + if not json_path.exists(): + return None + with open(json_path) as f: + d = json.load(f) + return np.array(d["q"]), np.array(d["Sq"]) + + +def compute_r_factor( + q_exp: np.ndarray, + sq_exp: np.ndarray, + q_calc: np.ndarray, + sq_calc: np.ndarray, +) -> float: + """ + Compute R-factor: sum|S_exp - S_calc| / sum|S_exp|. + + Parameters + ---------- + q_exp : np.ndarray + Experimental scattering vector values. + sq_exp : np.ndarray + Experimental structure factor values. + q_calc : np.ndarray + Calculated scattering vector values. + sq_calc : np.ndarray + Calculated structure factor values. + + Returns + ------- + float + R-factor value, or NaN if no overlapping valid data exists. + """ + valid_calc = ~np.isnan(sq_calc) + q_min_calc = float(q_calc[valid_calc].min()) + q_max_exp = float(q_exp.max()) + + f = interp1d( + q_calc[valid_calc], + sq_calc[valid_calc], + kind="linear", + bounds_error=False, + fill_value=np.nan, + ) + + mask = (q_exp >= q_min_calc) & (q_exp <= q_max_exp) + sq_interp = f(q_exp[mask]) + sq_e = sq_exp[mask] + + ok = ~np.isnan(sq_interp) + if not ok.any(): + return float("nan") + + return float(np.sum(np.abs(sq_e[ok] - sq_interp[ok])) / np.sum(np.abs(sq_e[ok]))) + + +def find_first_peak( + q: np.ndarray, sq: np.ndarray, q_min: float = 0.5, q_max: float = 2.0 +) -> float: + """ + Find position of first peak in S(q). + + Parameters + ---------- + q : np.ndarray + Scattering vector values. + sq : np.ndarray + Structure factor values. + q_min : float, optional + Lower bound of the search window in inverse angstroms (default 0.5). + q_max : float, optional + Upper bound of the search window in inverse angstroms (default 2.0). + + Returns + ------- + float + Position of the first peak, or NaN if fewer than 3 valid points + exist in the search window. + """ + mask = (q >= q_min) & (q <= q_max) & ~np.isnan(sq) + if mask.sum() < 3: + return float("nan") + q_sub = q[mask] + sq_sub = sq[mask] + return float(q_sub[np.argmax(sq_sub)]) + + +def build_sq_comparison_plot( + q_exp: np.ndarray, + sq_exp: np.ndarray, + model_data: dict[str, tuple[np.ndarray, np.ndarray]], +) -> go.Figure: + """ + Build S(q) overlay plot comparing computed and experimental data. + + Parameters + ---------- + q_exp : np.ndarray + Experimental scattering vector values. + sq_exp : np.ndarray + Experimental structure factor values. + model_data : dict of str to tuple of (np.ndarray, np.ndarray) + Mapping from model name to (q, S(q)) arrays. + + Returns + ------- + go.Figure + Plotly figure with experimental and computed S(q) traces. + """ + fig = go.Figure() + + fig.add_trace( + go.Scatter( + x=q_exp.tolist(), + y=sq_exp.tolist(), + mode="lines", + name="Exp. (Zhang et al. 2021)", + line={"color": "black", "width": 2}, + ) + ) + + for model, (q, sq) in model_data.items(): + fig.add_trace( + go.Scatter( + x=q.tolist(), + y=sq.tolist(), + mode="lines", + name=model, + opacity=0.8, + ) + ) + + fig.update_layout( + title="X-ray Structure Factor S(q) — LiTFSI/H2O (21 m)", + xaxis_title="q / Å⁻¹", + yaxis_title="S(q) (Faber-Ziman)", + legend={"x": 0.60, "y": 0.98}, + ) + return fig + + +# ============================================================================= +# Pytest fixtures +# ============================================================================= + + +@pytest.fixture +def density_results() -> dict[str, dict]: + """ + Load density results and write density plot JSONs. + + Returns + ------- + dict[str, dict] + Per-model density results keyed by registry model name. + """ + data = load_density_results() + if not data: + pytest.skip("No density data found") + + OUT_PATH.mkdir(parents=True, exist_ok=True) + + fig_bar = build_density_bar_chart(data) + with open(OUT_PATH / "figure_density_bar.json", "w") as f: + f.write(fig_bar.to_json()) + + fig_ts = build_density_timeseries(data) + with open(OUT_PATH / "figure_density_timeseries.json", "w") as f: + f.write(fig_ts.to_json()) + + return data + + +@pytest.fixture +def rdf_results() -> dict[str, dict]: + """ + Load RDF results and write RDF plot JSONs. + + Returns + ------- + dict[str, dict] + Per-model RDF results keyed by registry model name. + """ + data = load_rdf_results() + if not data: + pytest.skip("No RDF data found") + + OUT_PATH.mkdir(parents=True, exist_ok=True) + + fig_cn = build_cn_bar_chart(data) + with open(OUT_PATH / "figure_cn_bar.json", "w") as f: + f.write(fig_cn.to_json()) + + fig_gr = build_gr_plot(data) + with open(OUT_PATH / "figure_gr.json", "w") as f: + f.write(fig_gr.to_json()) + + return data + + +@pytest.fixture +def xray_sf_results() -> dict[str, dict]: + """ + Compute per-model R-factor and peak-position errors and write S(q) plot. + + Returns + ------- + dict[str, dict] + Per-model results keyed by registry model name, containing + ``r_factor``, ``peak_calc``, ``peak_exp``, ``peak_position_error``. + """ + q_exp, sq_exp = load_experimental_sq() + peak_exp = find_first_peak(q_exp, sq_exp) + + model_data: dict[str, tuple[np.ndarray, np.ndarray]] = {} + results: dict[str, dict] = {} + + for model in MODELS: + data = load_computed_sq(model) + if data is None: + continue + q_calc, sq_calc = data + model_data[model] = (q_calc, sq_calc) + + r_factor = compute_r_factor(q_exp, sq_exp, q_calc, sq_calc) + peak_calc = find_first_peak(q_calc, sq_calc) + peak_error = ( + abs(peak_calc - peak_exp) if not np.isnan(peak_calc) else float("nan") + ) + + results[model] = { + "r_factor": r_factor, + "peak_calc": peak_calc, + "peak_exp": peak_exp, + "peak_position_error": peak_error, + } + + if not results: + pytest.skip("No S(q) data found") + + OUT_PATH.mkdir(parents=True, exist_ok=True) + + fig = build_sq_comparison_plot(q_exp, sq_exp, model_data) + with open(OUT_PATH / "figure_xray_sq_comparison.json", "w") as f: + f.write(fig.to_json()) + + return results + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "litfsi_h2o_21m_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, +) +def metrics( + density_results: dict[str, dict], + rdf_results: dict[str, dict], + xray_sf_results: dict[str, dict], +) -> dict[str, dict]: + """ + Build per-metric dicts consumed by the ``build_table`` decorator. + + Parameters + ---------- + density_results + Per-model density results. + rdf_results + Per-model RDF / coordination-number results. + xray_sf_results + Per-model R-factor and peak-position results. + + Returns + ------- + dict[str, dict] + Metric names mapped to ``{model_name: value}`` dicts. The metric + order matches ``metrics.yml`` (Density Error, Density Error (%), + CN Li-O_water Error, CN Li-O_TFSI Error, CN Li-O_total Error, + S(q) R-factor, First Peak Position Error). + """ + return { + "Density Error": { + model: d["rho_abs_error"] for model, d in density_results.items() + }, + "Density Error (%)": { + model: abs(d["rho_error_pct"]) for model, d in density_results.items() + }, + "CN Li-O_water Error": { + model: d["err_water"] for model, d in rdf_results.items() + }, + "CN Li-O_TFSI Error": { + model: d["err_tfsi"] for model, d in rdf_results.items() + }, + "CN Li-O_total Error": { + model: d["err_total"] for model, d in rdf_results.items() + }, + "S(q) R-factor": {model: r["r_factor"] for model, r in xray_sf_results.items()}, + "First Peak Position Error": { + model: r["peak_position_error"] for model, r in xray_sf_results.items() + }, + } + + +def test_litfsi_h2o_21m(metrics: dict[str, dict]) -> None: + """ + Run the consolidated WiSE LiTFSI/H2O benchmark. + + Parameters + ---------- + metrics + Per-metric values for all models. + """ + assert metrics, "No models produced data for the WiSE LiTFSI/H2O benchmark" diff --git a/ml_peg/analysis/electrolytes/litfsi_h2o_21m/metrics.yml b/ml_peg/analysis/electrolytes/litfsi_h2o_21m/metrics.yml new file mode 100644 index 000000000..d1f786af4 --- /dev/null +++ b/ml_peg/analysis/electrolytes/litfsi_h2o_21m/metrics.yml @@ -0,0 +1,50 @@ +metrics: + Density Error: + good: 0.0 + bad: 0.17 + unit: "g/cm3" + tooltip: "Absolute error in NPT density vs experimental (1.7126 g/cm3, Gilbert et al. JCED 2017)" + level_of_theory: Experimental (Gilbert et al., J. Chem. Eng. Data 62, 2056, 2017) + weight: 1 + Density Error (%): + good: 0.0 + bad: 10.0 + unit: "%" + tooltip: "Percent error in NPT density vs experimental (1.7126 g/cm3, Gilbert et al. JCED 2017)" + level_of_theory: Experimental (Gilbert et al., J. Chem. Eng. Data 62, 2056, 2017) + weight: 0 + CN Li-O_water Error: + good: 0.15 + bad: 0.5 + unit: "-" + tooltip: "Absolute error in Li-O_water coordination number vs theoretical reference at 21 m (CN=2.4: r2SCAN AIMD 2.23 at 20 ps; MACE potentials trained on the r2SCAN dataset converge to 2.35-2.39 over >2 ns). The good threshold (0.15) equals the reference uncertainty (AIMD vs converged spread). The experimental split (Watanabe et al. JPCB 2021) is at ~18.5 m and not used for the partitioning." + level_of_theory: Theory (r2SCAN AIMD + r2SCAN-trained MACE, 21 m) + weight: 0.25 + CN Li-O_TFSI Error: + good: 0.05 + bad: 0.5 + unit: "-" + tooltip: "Absolute error in Li-O_TFSI coordination number vs theoretical reference at 21 m (CN=1.9: r2SCAN AIMD 2.02 at 20 ps; MACE potentials trained on the r2SCAN dataset converge to 1.89-1.91 over >2 ns). The good threshold (0.05) equals the reference uncertainty (FT/TfS agreement). The experimental split (Watanabe et al. JPCB 2021) is at ~18.5 m and not used for the partitioning." + level_of_theory: Theory (r2SCAN AIMD + r2SCAN-trained MACE, 21 m) + weight: 0.25 + CN Li-O_total Error: + good: 0.10 + bad: 0.5 + unit: "-" + tooltip: "Absolute error in total Li-O coordination number vs reference CN=4.3 (r2SCAN AIMD 4.31; converged r2SCAN-trained MACE 4.27-4.29; consistent with the experimental 4.21 at ~18.5 m, Watanabe et al. JPCB 2021, neutron diffraction with isotopic substitution). The good threshold (0.10) equals the theory-experiment spread." + level_of_theory: Theory (r2SCAN, 21 m), consistent with experiment at 18.5 m + weight: 0.5 + S(q) R-factor: + good: 0.0 + bad: 0.5 + unit: "-" + tooltip: "R-factor between computed and experimental X-ray S(q): sum|S_exp - S_calc| / sum|S_exp|" + level_of_theory: Experimental (SAXS, Zhang et al., J. Phys. Chem. B 125, 4501, 2021) + weight: 1 + First Peak Position Error: + good: 0.0 + bad: 0.3 + unit: "1/A" + tooltip: "Absolute error in position of first S(q) peak (~1.0 A-1)" + level_of_theory: Experimental (SAXS, Zhang et al., J. Phys. Chem. B 125, 4501, 2021) + weight: 0 diff --git a/ml_peg/app/electrolytes/electrolytes.yml b/ml_peg/app/electrolytes/electrolytes.yml new file mode 100644 index 000000000..437539c78 --- /dev/null +++ b/ml_peg/app/electrolytes/electrolytes.yml @@ -0,0 +1,5 @@ +title: Electrolytes +description: Structural, thermodynamic and transport properties of liquid and solid-state electrolytes +weight: 0 +benchmark_weights: + WiSE LiTFSI/H2O 21 m: 1.0 diff --git a/ml_peg/app/electrolytes/litfsi_h2o_21m/app_litfsi_h2o_21m.py b/ml_peg/app/electrolytes/litfsi_h2o_21m/app_litfsi_h2o_21m.py new file mode 100644 index 000000000..bc31bd086 --- /dev/null +++ b/ml_peg/app/electrolytes/litfsi_h2o_21m/app_litfsi_h2o_21m.py @@ -0,0 +1,87 @@ +"""Dash app for the consolidated WiSE 21 m LiTFSI/H2O electrolyte benchmark.""" + +from __future__ import annotations + +from pathlib import Path + +from dash import html + +from ml_peg.app.base_app import BaseApp +from ml_peg.app.utils.build_callbacks import plot_from_table_column +from ml_peg.app.utils.load import read_plot + +BENCHMARK_NAME = "electrolytes_litfsi_h2o_21m" +APP_ROOT = Path(__file__).resolve().parents[2] +DATA_PATH = APP_ROOT / "data" / "electrolytes" / "litfsi_h2o_21m" +DOCS_URL = ( + "https://ddmms.github.io/ml-peg/user_guide/benchmarks/" + "electrolytes.html#litfsi-h2o-21-m" +) + + +class LitfsiH2O21mApp(BaseApp): + """Dash app for the consolidated WiSE LiTFSI/H2O 21 m benchmark.""" + + def register_callbacks(self) -> None: + """Register interactive plot callbacks for all metrics.""" + density_bar = read_plot( + DATA_PATH / "figure_density_bar.json", + id=f"{BENCHMARK_NAME}-density-bar", + ) + cn_bar = read_plot( + DATA_PATH / "figure_cn_bar.json", + id=f"{BENCHMARK_NAME}-cn-bar", + ) + gr_plot = read_plot( + DATA_PATH / "figure_gr.json", + id=f"{BENCHMARK_NAME}-gr", + ) + sq_plot = read_plot( + DATA_PATH / "figure_xray_sq_comparison.json", + id=f"{BENCHMARK_NAME}-sq", + ) + + plot_from_table_column( + table_id=self.table_id, + plot_id=f"{BENCHMARK_NAME}-figure-placeholder", + column_to_plot={ + "Density Error": density_bar, + "Density Error (%)": density_bar, + "CN Li-O_water Error": cn_bar, + "CN Li-O_TFSI Error": cn_bar, + "S(q) R-factor": sq_plot, + "First Peak Position Error": sq_plot, + "Score": gr_plot, + }, + ) + + +def get_app() -> LitfsiH2O21mApp: + """ + Return the configured WiSE LiTFSI/H2O 21 m benchmark app. + + Returns + ------- + LitfsiH2O21mApp + Configured app instance. + """ + return LitfsiH2O21mApp( + name="WiSE LiTFSI/H2O 21 m", + description=( + "Consolidated structural and thermodynamic benchmark for the " + "21 m LiTFSI/H2O 'water-in-salt' electrolyte: NPT density vs " + "experiment (Gilbert 2017), Li-O coordination numbers from RDFs " + "(Watanabe 2021), and X-ray structure factor S(q) (Zhang 2021). " + "Trajectories produced with LAMMPS+symmetrix on Adastra (MI250X)." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "litfsi_h2o_21m_metrics_table.json", + extra_components=[ + html.Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), + ], + ) + + +if __name__ == "__main__": + app = get_app() + app.run(port=8060, debug=True) diff --git a/ml_peg/calcs/electrolytes/litfsi_h2o_21m/calc_litfsi_h2o_21m.py b/ml_peg/calcs/electrolytes/litfsi_h2o_21m/calc_litfsi_h2o_21m.py new file mode 100644 index 000000000..d583f3b35 --- /dev/null +++ b/ml_peg/calcs/electrolytes/litfsi_h2o_21m/calc_litfsi_h2o_21m.py @@ -0,0 +1,959 @@ +""" +Consolidated WiSE 21 m LiTFSI/H2O electrolyte benchmark. + +The benchmark has two halves, which can be run together or independently. + +**Reference MD** (``test_reference_md``, marked ``very_slow``) reproduces the +protocol that generated the reference data, and writes its products into +``DATA_ROOT//``: + +* ``p64_w170`` (1534 atoms, 27.4938 A cubic): Min -> NVT 50 ps equilibration + -> NVT 50 ps production, held at the experimental volume throughout. No NPT: + S(q) and the coordination numbers are compared at the reference density + rather than at each model's own. Produces ``nvt_trajectory.extxyz``. +* ``p16_w42`` (382 atoms): Min -> NVT 50 ps -> NPT 200 ps. NPT is used only on + this smaller cell, where the density converges far more cheaply. Produces + ``density.json``, averaged over the last 150 ps. + +**Extraction** (the remaining tests, fast) reads those files back and computes +three observables for each registered MLIP model: + +* Density from the NPT run (p16_w42). +* Li-O coordination numbers from the NVT trajectory via radial distribution + functions g(r) for Li-O_water (O bonded to H, d_OH < 1.25 A) and + Li-O_TFSI (O bonded to S, d_OS < 1.75 A); CN integrated to first minimum + R_CUT = 2.83 A from the r2SCAN AIMD reference. +* X-ray structure factor S(q) computed via dynasor in TRAVIS Faber-Ziman + convention with Cromer-Mann 4-Gaussian form factors (including H) and + Savitzky-Golay smoothing (window=5, order=3, dq=0.02 A^-1; physical width 0.10 A^-1). + +The published reference data were produced with LAMMPS + symmetrix on Adastra +(MI250X) following the same protocol; ``test_reference_md`` is the janus-core +recast of it, and lets the data be regenerated for any registered model. +Integrators are matched where janus-core exposes a choice: NVT uses the +Nose-Hoover chain (``NVT_NH``, equivalent to LAMMPS ``fix nvt``) and NPT the +Martyna-Tobias-Klein chain (``NPT_MTK``, the formulation used by LAMMPS +``fix npt``; plain ``NPT`` in janus-core is Melchionna and would not match). + +Reference values: + +* Density: 1.7126 g/cm3 (experimental; Gilbert et al., J. Chem. Eng. Data + 62, 2056, 2017). +* Li-O CN: 2.4 (water), 1.9 (TFSI), 4.3 (total). Theoretical references at + 21 m: r2SCAN AIMD (Li128, 20 ps) gives 2.23/2.02/4.31; MACE potentials + trained on the same r2SCAN dataset converge to 2.35-2.39/1.89-1.91/ + 4.27-4.29 over >2 ns. The total is consistent with the experimental 4.21 + at ~18.5 m (Watanabe et al., JPCB 125, 7477, 2021, neutron diffraction + with isotopic substitution); the experimental water/TFSI split at that + concentration is not used, as it disagrees systematically with all MD + partitionings and is measured at a different composition. +* S(q): SAXS (Zhang et al., J. Phys. Chem. B 125, 4501, 2021). + +Data live under ``DATA_ROOT//``, where ``DATA_ROOT`` +defaults to ``./data`` next to this script and is overridable via the +``ML_PEG_WISE_LITFSI_H2O_21M_DATA_ROOT`` environment variable. The starting +structures are expected in ``DATA_ROOT/structures/``. +""" + +from __future__ import annotations + +from copy import copy +import json +import os +from pathlib import Path +from typing import Any +import warnings + +from ase import Atoms +from ase.geometry import get_distances +from ase.io import iread +from ase.io import read as ase_read +from ase.io import write as ase_write +import numpy as np +import pytest +from scipy.signal import savgol_filter + +from ml_peg.models import current_models +from ml_peg.models.get_models import load_models + +warnings.filterwarnings("ignore") + +# --- Configuration ----------------------------------------------------------- + +MODELS = load_models(current_models) + +DATA_ROOT = Path( + os.environ.get( + "ML_PEG_WISE_LITFSI_H2O_21M_DATA_ROOT", + Path(__file__).parent / "data", + ) +) +OUT_PATH = Path(__file__).parent / "outputs" + +# --- Density references ------------------------------------------------------ + +RHO_EXP = 1.7126 # g/cm3 (Gilbert et al., JCED 2017) + +# --- RDF parameters ---------------------------------------------------------- + +R_MAX = 6.0 # A +DR = 0.02 # A +R_CUT_COORD = 2.83 # A — first minimum of Li-O g(r) from r2SCAN AIMD +D_OH_CUT = 1.25 # A — O-H bond cutoff for water identification +D_OS_CUT = 1.75 # A — O-S bond cutoff for TFSI identification + +# --- S(q) parameters (TRAVIS-matching) --------------------------------------- + +Q_MAX = 13.0 # A^-1 +Q_MIN = 0.5 # A^-1 +MAX_QPOINTS = 50000 +DQ_BIN = 0.02 # A^-1 +SAVGOL_WINDOW = 5 +SAVGOL_ORDER = 3 + +# LAMMPS atom-type → element fallback for raw dump files +TYPE_TO_ELEMENT = {1: "Li", 2: "C", 3: "F", 4: "S", 5: "N", 6: "O", 7: "H"} + +# System composition (p64_w170: 64 LiTFSI + 170 H2O) +COMPOSITION = {"Li": 64, "C": 128, "F": 384, "S": 128, "N": 64, "O": 426, "H": 340} +N_ATOMS = sum(COMPOSITION.values()) # 1534 +CONC = {k: v / N_ATOMS for k, v in COMPOSITION.items()} + +# Cromer-Mann 4-Gaussian X-ray form factor parameters (International Tables) +TRAVIS_FF = { + "S": { + "a": [6.905, 5.203, 1.438, 1.586], + "b": [1.468, 22.215, 0.254, 56.172], + "c": 0.867, + }, + "F": { + "a": [3.539, 2.641, 1.517, 1.024], + "b": [10.283, 4.294, 0.262, 26.148], + "c": 0.278, + }, + "O": { + "a": [3.049, 2.287, 1.546, 0.867], + "b": [13.277, 5.701, 0.324, 32.909], + "c": 0.251, + }, + "N": { + "a": [12.213, 3.132, 2.013, 1.166], + "b": [0.006, 9.893, 28.997, 0.583], + "c": -11.529, + }, + "C": { + "a": [2.310, 1.020, 1.589, 0.865], + "b": [20.844, 10.208, 0.569, 51.651], + "c": 0.216, + }, + "Li": { + "a": [1.128, 0.751, 0.618, 0.465], + "b": [3.955, 1.052, 85.391, 168.261], + "c": 0.038, + }, + "H": { + "a": [0.493, 0.323, 0.140, 0.041], + "b": [10.511, 26.126, 3.142, 57.800], + "c": 0.003, + }, +} + +# --- Reference MD protocol --------------------------------------------------- + +TEMPERATURE_K = 298.15 +PRESSURE_BAR = 1.01325 # 1 atm +PRESSURE_GPA = PRESSURE_BAR * 1e-4 # janus-core takes pressure in GPa +TIMESTEP_FS = 0.5 + +# Nose-Hoover damping times (LAMMPS TDAMP/PDAMP: 100*dt and 1000*dt) +THERMOSTAT_TIME_FS = 50.0 +BAROSTAT_TIME_FS = 500.0 +NH_CHAIN = 3 # LAMMPS default for both thermostat and barostat sub-chains + +# Step counts at 0.5 fs/step +NVT_EQUIL_STEPS = 100_000 # 50 ps, both cells +NVT_PROD_STEPS = 100_000 # 50 ps, p64_w170 only (S(q), RDF) +NPT_PROD_STEPS = 400_000 # 200 ps, p16_w42 only (density) + +# IO cadence (LAMMPS THERMO_EVERY / DUMP_EVERY) +STATS_EVERY = 100 # 0.05 ps +TRAJ_EVERY = 200 # 0.1 ps -> 501 frames over the 50 ps production run + +# Minimization (LAMMPS: min_style cg; minimize 1.0e-6 0.2 2000 20000) +MIN_FMAX = 0.2 # eV/A +MIN_STEPS = 2000 + +SEED = 42 + +# Density is averaged over the last 150 ps of the 200 ps NPT run. +DENSITY_WINDOW_PS = (50.0, 200.0) + +STRUCTURE_DIR = DATA_ROOT / "structures" +CELL_P64 = "p64_w170" +CELL_P16 = "p16_w42" + + +# ============================================================================= +# Reference MD (slow: regenerates the data consumed by the extraction below) +# ============================================================================= + + +def load_initial_structure(cell: str) -> Atoms: + """ + Load a packed starting configuration at the experimental density. + + Parameters + ---------- + cell : str + Cell name, ``"p64_w170"`` (1534 atoms) or ``"p16_w42"`` (382 atoms). + + Returns + ------- + Atoms + ASE Atoms at rho = 1.7126 g/cm3. + + Raises + ------ + FileNotFoundError + If the structure is not available locally. + """ + path = STRUCTURE_DIR / f"{cell}_initial.xyz" + if not path.exists(): + raise FileNotFoundError( + f"Starting structure for {cell} not found at {path}. " + "Both cells (p64_w170, p16_w42) are needed to run the reference MD." + ) + struct = ase_read(path) + + # Both cells are neutral, closed-shell (equal numbers of Li+ and TFSI-). + # Models trained on molecular data, such as orb-v3-consv-omol, refuse to + # evaluate a structure that does not carry these keys. + struct.info.setdefault("charge", 0) + struct.info.setdefault("spin", 1) + + return struct + + +def _run_stage( + ensemble_cls, struct: Atoms, *, steps: int, file_prefix: Path, **kwargs +) -> Atoms: + """ + Run one MD stage, skipping it if it already completed. + + janus-core writes ``{file_prefix}-final.extxyz`` once a stage reaches its + last step, so re-running the test picks up from the first stage that did + not finish rather than repeating the whole protocol. + + Parameters + ---------- + ensemble_cls + The janus-core ensemble class, e.g. ``NVT_NH`` or ``NPT_MTK``. + struct : Atoms + Structure to propagate. + steps : int + Number of steps for this stage. + file_prefix : Path + Prefix for this stage's output files. + **kwargs + Further arguments forwarded to the ensemble class. + + Returns + ------- + Atoms + The propagated structure, to hand on to the next stage. + """ + final_file = Path(f"{file_prefix}-final.extxyz") + if final_file.exists(): + done = ase_read(final_file) + done.calc = struct.calc + return done + + md = ensemble_cls( + struct=struct, + steps=steps, + temp=TEMPERATURE_K, + timestep=TIMESTEP_FS, + thermostat_time=THERMOSTAT_TIME_FS, + stats_every=STATS_EVERY, + seed=SEED, + file_prefix=file_prefix, + **kwargs, + ) + md.run() + + # Return the structure janus-core actually propagated rather than the one + # passed in, which it is free to rebind. + return md.struct + + +def _write_density_json(model_name: str, stats_file: Path, out_file: Path) -> dict: + """ + Average the NPT density over ``DENSITY_WINDOW_PS`` and save it. + + Parameters + ---------- + model_name : str + Name of the MLIP model in the registry. + stats_file : Path + The janus-core stats file of the NPT stage. + out_file : Path + Where to write ``density.json``. + + Returns + ------- + dict + The density summary that was written. + """ + from janus_core.helpers.stats import Stats + + stats = Stats(stats_file) + time_ps = np.asarray(stats["Time"], dtype=float) / 1000.0 # fs -> ps + rho = np.asarray(stats["Density"], dtype=float) + + lo, hi = DENSITY_WINDOW_PS + window = rho[(time_ps >= lo) & (time_ps <= hi)] + if window.size == 0: + raise RuntimeError( + f"No NPT samples in {lo}-{hi} ps for {model_name}; the run is too short." + ) + + summary = { + "model": model_name, + "cell": CELL_P16, + "rho_exp": RHO_EXP, + "rho_mean": float(window.mean()), + "rho_std": float(window.std()), + "rho_error_pct": float(100 * (window.mean() - RHO_EXP) / RHO_EXP), + "rho_abs_error": float(abs(window.mean() - RHO_EXP)), + "n_samples": int(window.size), + "time_range_ps": [lo, hi], + "time_full": time_ps.tolist(), + "density_full": rho.tolist(), + } + out_file.parent.mkdir(parents=True, exist_ok=True) + out_file.write_text(json.dumps(summary, indent=2)) + return summary + + +def run_reference_md(model_name: str, model: Any) -> None: + """ + Run the reference protocol for one model and write its data products. + + Produces ``nvt_trajectory.extxyz`` (p64_w170) and ``density.json`` + (p16_w42) under ``DATA_ROOT//``, which the extraction tests + then read back. + + Parameters + ---------- + model_name : str + Registry name of the MLIP model. + model : Any + Model object from :func:`load_models`. + """ + # Imported lazily: janus-core pulls in torch, which is not needed by the + # extraction tests. + from janus_core.calculations.geom_opt import GeomOpt + from janus_core.calculations.md import NPT_MTK, NVT_NH + + calc = model.get_calculator(precision="high") + + data_dir = DATA_ROOT / model_name + work_dir = OUT_PATH / model_name / "reference_md" + work_dir.mkdir(parents=True, exist_ok=True) + + for cell in (CELL_P64, CELL_P16): + struct = load_initial_structure(cell) + struct.calc = copy(calc) + + # filter_class=None relaxes the atoms only: LAMMPS `minimize` leaves the + # cell alone, and both cells must start from the experimental volume. + GeomOpt( + struct=struct, + fmax=MIN_FMAX, + optimizer="FIRE", + steps=MIN_STEPS, + filter_class=None, + write_traj=False, + file_prefix=work_dir / f"{cell}-minimize", + ).run() + + # The equilibration trajectory is written but not analysed, as in the + # LAMMPS pipeline. (traj_every=0 is not a way to disable it: ASE reads a + # non-positive interval as "call once at step 0", which would divide by + # zero in janus-core's trajectory writer.) + struct = _run_stage( + NVT_NH, + struct, + steps=NVT_EQUIL_STEPS, + file_prefix=work_dir / f"{cell}-nvt_equil", + traj_every=TRAJ_EVERY, + ) + + if cell == CELL_P64: + # Production at the experimental volume: this is what S(q) and the + # RDF are computed from. + prefix = work_dir / f"{cell}-nvt_prod" + _run_stage( + NVT_NH, + struct, + steps=NVT_PROD_STEPS, + file_prefix=prefix, + traj_every=TRAJ_EVERY, + ) + traj = ase_read(f"{prefix}-traj.extxyz", index=":") + data_dir.mkdir(parents=True, exist_ok=True) + ase_write(data_dir / "nvt_trajectory.extxyz", traj, format="extxyz") + else: + # NPT only on the small cell, for the density. + prefix = work_dir / f"{cell}-npt_prod" + _run_stage( + NPT_MTK, + struct, + steps=NPT_PROD_STEPS, + file_prefix=prefix, + pressure=PRESSURE_GPA, + barostat_time=BAROSTAT_TIME_FS, + thermostat_chain=NH_CHAIN, + barostat_chain=NH_CHAIN, + traj_every=TRAJ_EVERY, + ) + _write_density_json( + model_name, + Path(f"{prefix}-stats.dat"), + data_dir / "density.json", + ) + + +@pytest.mark.very_slow +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_reference_md(mlip: tuple[str, Any]) -> None: + """ + Regenerate the benchmark data for one model with the reference protocol. + + Around 300 ps of cumulative MD on 1534 and 382 atoms: expect on the order + of a day of GPU time per model. + + Parameters + ---------- + mlip : tuple[str, Any] + ``(model_name, model)`` pair from the ml-peg registry. + """ + model_name, model = mlip + run_reference_md(model_name, model) + + assert (DATA_ROOT / model_name / "nvt_trajectory.extxyz").exists() + assert (DATA_ROOT / model_name / "density.json").exists() + + +# ============================================================================= +# Density extraction +# ============================================================================= + + +def _density_source_path(model_name: str) -> Path: + """ + Locate the pre-computed density JSON for a model. + + Parameters + ---------- + model_name : str + Name of the MLIP model in the registry. + + Returns + ------- + Path + Expected path to ``density.json`` under ``DATA_ROOT``. + """ + return DATA_ROOT / model_name / "density.json" + + +# ============================================================================= +# RDF computation +# ============================================================================= + + +def identify_o_types(atoms) -> tuple[np.ndarray, np.ndarray]: + """ + Return indices of O_water and O_TFSI from a single ASE Atoms frame. + + Parameters + ---------- + atoms : ase.Atoms + A single frame from the trajectory. + + Returns + ------- + o_water : np.ndarray + Indices of oxygen atoms bonded to hydrogen (water oxygens). + o_tfsi : np.ndarray + Indices of oxygen atoms bonded to sulfur (TFSI oxygens). + """ + syms = np.array(atoms.get_chemical_symbols()) + pos = atoms.get_positions() + cell = atoms.get_cell() + pbc = atoms.get_pbc() + + o_idx = np.where(syms == "O")[0] + h_idx = np.where(syms == "H")[0] + s_idx = np.where(syms == "S")[0] + + o_water, o_tfsi = [], [] + for o in o_idx: + _, d_oh = get_distances(pos[o : o + 1], pos[h_idx], cell=cell, pbc=pbc) + _, d_os = get_distances(pos[o : o + 1], pos[s_idx], cell=cell, pbc=pbc) + if d_oh.min() < D_OH_CUT: + o_water.append(o) + elif d_os.min() < D_OS_CUT: + o_tfsi.append(o) + + return np.array(o_water), np.array(o_tfsi) + + +def compute_rdf( + traj_path: Path, + o_water_idx: np.ndarray, + o_tfsi_idx: np.ndarray, + r_max: float = R_MAX, + dr: float = DR, + skip_frames: int = 0, +) -> dict: + """ + Compute Li-O RDFs from an extxyz trajectory. + + Parameters + ---------- + traj_path : Path + Path to .extxyz trajectory file. + o_water_idx : np.ndarray + Atom indices of O_water (from first frame, fixed). + o_tfsi_idx : np.ndarray + Atom indices of O_TFSI (from first frame, fixed). + r_max : float + Maximum distance for RDF. + dr : float + Bin width. + skip_frames : int + Number of initial frames to skip (equilibration). Trajectories are + already pre-equilibrated (50–100 ps window). + + Returns + ------- + dict + Dictionary with keys ``r``, ``gr_LiO_total``, ``gr_LiO_water``, + ``gr_LiO_TFSI``, ``coord_LiO_total``, ``coord_LiO_water``, + ``coord_LiO_TFSI``, ``n_li``, ``n_O_water``, ``n_O_TFSI``, + ``n_frames_used``, and ``r_cut_coord``. + """ + bins = np.arange(0, r_max + dr, dr) + r_centers = 0.5 * (bins[:-1] + bins[1:]) + n_bins = len(r_centers) + + hist_total = np.zeros(n_bins) + hist_water = np.zeros(n_bins) + hist_tfsi = np.zeros(n_bins) + + n_frames = 0 + n_li = None + volume = None + o_all_idx = np.concatenate([o_water_idx, o_tfsi_idx]) + + for frame_idx, atoms in enumerate( + iread(str(traj_path), format="extxyz", index=":") + ): + if frame_idx < skip_frames: + continue + + syms = np.array(atoms.get_chemical_symbols()) + pos = atoms.get_positions() + cell = atoms.get_cell() + pbc = atoms.get_pbc() + + li_idx = np.where(syms == "Li")[0] + if n_li is None: + n_li = len(li_idx) + if volume is None: + volume = atoms.get_volume() + + pos_li = pos[li_idx] + + for o_set, hist in [ + (o_all_idx, hist_total), + (o_water_idx, hist_water), + (o_tfsi_idx, hist_tfsi), + ]: + pos_o = pos[o_set] + _, dists = get_distances(pos_li, pos_o, cell=cell, pbc=pbc) + dists_flat = dists.ravel() + dists_flat = dists_flat[dists_flat < r_max] + h, _ = np.histogram(dists_flat, bins=bins) + hist += h + + n_frames += 1 + + if n_frames == 0 or n_li is None: + raise RuntimeError(f"No frames processed from {traj_path}") + + def normalize(hist, n_central, n_neighbor): + """ + Normalize histogram to g(r). + + Parameters + ---------- + hist : np.ndarray + Raw pair-distance histogram. + n_central : int + Number of central atoms (Li). + n_neighbor : int + Number of neighbor atoms (O subset). + + Returns + ------- + np.ndarray + Normalized radial distribution function. + """ + shell_vol = (4.0 / 3.0) * np.pi * (bins[1:] ** 3 - bins[:-1] ** 3) + rho = n_neighbor / volume + norm = n_central * n_frames * rho * shell_vol + return hist / norm + + n_o_total = len(o_all_idx) + n_o_water = len(o_water_idx) + n_o_tfsi = len(o_tfsi_idx) + + gr_total = normalize(hist_total, n_li, n_o_total) + gr_water = normalize(hist_water, n_li, n_o_water) + gr_tfsi = normalize(hist_tfsi, n_li, n_o_tfsi) + + def coord_number(gr, n_neighbor): + """ + Compute coordination number from g(r). + + Parameters + ---------- + gr : np.ndarray + Radial distribution function. + n_neighbor : int + Number of neighbor atoms (O subset). + + Returns + ------- + float + Integrated coordination number up to R_CUT_COORD. + """ + rho = n_neighbor / volume + integrand = 4.0 * np.pi * rho * gr * r_centers**2 * dr + mask = r_centers <= R_CUT_COORD + return float(np.sum(integrand[mask])) + + return { + "r": r_centers.tolist(), + "gr_LiO_total": gr_total.tolist(), + "gr_LiO_water": gr_water.tolist(), + "gr_LiO_TFSI": gr_tfsi.tolist(), + "coord_LiO_total": coord_number(gr_total, n_o_total), + "coord_LiO_water": coord_number(gr_water, n_o_water), + "coord_LiO_TFSI": coord_number(gr_tfsi, n_o_tfsi), + "n_li": n_li, + "n_O_water": n_o_water, + "n_O_TFSI": n_o_tfsi, + "n_frames_used": n_frames, + "r_cut_coord": R_CUT_COORD, + } + + +# ============================================================================= +# X-ray S(q) computation +# ============================================================================= + + +def compute_fq_travis(elem: str, q_arr: np.ndarray) -> np.ndarray: + """ + Compute X-ray form factor f(q) using TRAVIS 4-Gaussian (Cromer-Mann) params. + + Parameters + ---------- + elem : str + Chemical element symbol (e.g. ``"Li"``, ``"O"``). + q_arr : np.ndarray + Array of q values in inverse angstroms. + + Returns + ------- + np.ndarray + Form factor values evaluated at each q. + """ + ff = TRAVIS_FF[elem] + s2 = (q_arr / (4 * np.pi)) ** 2 + return sum(ff["a"][i] * np.exp(-ff["b"][i] * s2) for i in range(4)) + ff["c"] + + +def compute_sq_travis_style(traj_path: Path) -> dict: + """ + Compute S(q) in TRAVIS Faber-Ziman convention using dynasor. + + Steps: + + 1. Read trajectory with dynasor (ASE for .extxyz; MDAnalysis for raw dump). + 2. Compute partial S_ab(q) on a fine spherical q-grid. + 3. Apply TRAVIS Cromer-Mann form factors: I_xray = sum f_a * f_b * S_ab. + 4. Spherical binning at dq = 0.02 A^-1. + 5. Faber-Ziman normalization with Laue term: S_FZ = I/^2 - /^2 + 1. + 6. Savitzky-Golay smoothing. + + Parameters + ---------- + traj_path : Path + Path to the trajectory file (.extxyz or LAMMPS dump). + + Returns + ------- + dict + Dictionary with keys ``q``, ``Sq``, ``n_qpoints``, + ``n_qpoints_used``, ``cell``, ``atom_types``, + ``particle_counts``, and ``params``. + """ + from dynasor import ( + Trajectory, + compute_static_structure_factors, + get_spherical_qpoints, + ) + + traj_str = str(traj_path) + is_extxyz = traj_str.endswith(".extxyz") or traj_str.endswith(".xyz") + + if is_extxyz: + first_frame = ase_read(traj_str, index=0) + symbols = first_frame.get_chemical_symbols() + atomic_indices = {} + for i, sym in enumerate(symbols): + atomic_indices.setdefault(sym, []).append(i) + + traj = Trajectory( + traj_str, + trajectory_format="ase", + atomic_indices=atomic_indices, + ) + else: + import MDAnalysis as mda # noqa: N813 + + u = mda.Universe(traj_str, format="LAMMPSDUMP") + types = u.atoms.types + atomic_indices = {} + for t, elem in TYPE_TO_ELEMENT.items(): + mask = types == str(t) + idx = np.where(mask)[0].tolist() + if idx: + atomic_indices[elem] = idx + + traj = Trajectory( + traj_str, + trajectory_format="lammps_mdanalysis", + atomic_indices=atomic_indices, + ) + + q_points = get_spherical_qpoints(traj.cell, q_max=Q_MAX, max_points=MAX_QPOINTS) + sample = compute_static_structure_factors(traj, q_points) + + q_norms = np.linalg.norm(sample.q_points, axis=1) + atom_types = list(sample.particle_counts.keys()) + + ff_at_q = {at: compute_fq_travis(at, q_norms) for at in atom_types} + i_xray = np.zeros(len(q_norms)) + for s1, s2 in sample.pairs: + sq_ab = sample[f"Sq_{s1}_{s2}"].flatten() + i_xray += ff_at_q[s1] * ff_at_q[s2] * sq_ab + + q_bins = np.arange(0.0, Q_MAX + DQ_BIN, DQ_BIN) + q_centers = 0.5 * (q_bins[:-1] + q_bins[1:]) + i_xray_binned = np.full(len(q_centers), np.nan) + counts = np.zeros(len(q_centers), dtype=int) + + for i in range(len(q_centers)): + mask = (q_norms >= q_bins[i]) & (q_norms < q_bins[i + 1]) + n = mask.sum() + if n > 0: + i_xray_binned[i] = np.mean(i_xray[mask]) + counts[i] = n + + f_avg = np.zeros(len(q_centers)) + f2_avg = np.zeros(len(q_centers)) + for elem, c in CONC.items(): + fq = compute_fq_travis(elem, q_centers) + f_avg += c * fq + f2_avg += c * fq**2 + f_avg_sq = f_avg**2 + + sq_fz = np.where( + f_avg_sq > 0, + i_xray_binned / f_avg_sq - f2_avg / f_avg_sq + 1.0, + np.nan, + ) + + valid = ~np.isnan(sq_fz) & (q_centers >= 0.3) & (q_centers <= Q_MAX) + q_v = q_centers[valid] + sq_v = sq_fz[valid] + if len(sq_v) > SAVGOL_WINDOW: + sq_smooth = savgol_filter(sq_v, SAVGOL_WINDOW, SAVGOL_ORDER) + else: + sq_smooth = sq_v + + rmask = (q_v >= Q_MIN) & (q_v <= 12.0) + + return { + "q": q_v[rmask].tolist(), + "Sq": sq_smooth[rmask].tolist(), + "n_qpoints": len(q_points), + "n_qpoints_used": len(sample.q_points), + "cell": traj.cell.tolist(), + "atom_types": traj.atom_types, + "particle_counts": {k: int(v) for k, v in sample.particle_counts.items()}, + "params": { + "q_max": Q_MAX, + "q_min": Q_MIN, + "dq_bin": DQ_BIN, + "max_qpoints": MAX_QPOINTS, + "savgol_window": SAVGOL_WINDOW, + "savgol_order": SAVGOL_ORDER, + "form_factors": "cromer-mann-4gaussian", + "normalization": "faber-ziman", + }, + } + + +def find_trajectory(model_name: str) -> Path | None: + """ + Find NVT trajectory for a model. + + Prefers the converted ``.extxyz`` and falls back to the raw LAMMPS dump. + + Parameters + ---------- + model_name : str + Name of the MLIP model. + + Returns + ------- + Path or None + Path to the trajectory file, or ``None`` if not found. + """ + extxyz = DATA_ROOT / model_name / "nvt_trajectory.extxyz" + if extxyz.exists(): + return extxyz + + lammpstrj = DATA_ROOT / model_name / "nvt_trajectory.lammpstrj" + if lammpstrj.exists(): + return lammpstrj + + return None + + +# ============================================================================= +# Pytest interface (ml-peg convention) +# ============================================================================= + + +@pytest.mark.parametrize("model_name", MODELS) +def test_extract_density(model_name: str) -> None: + """ + Extract and save NPT density data for one model. + + Parameters + ---------- + model_name : str + Name of the MLIP model in the registry. + """ + src = _density_source_path(model_name) + if not src.exists(): + pytest.skip( + f"No density data for {model_name} at {src}. Fetch the reference " + "data or generate it with test_reference_md (-m very_slow)." + ) + + with open(src) as f: + result = json.load(f) + + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_dir / "density.json", "w") as f: + json.dump(result, f, indent=2) + + assert result["rho_mean"] > 0, f"Negative density for {model_name}" + assert abs(result["rho_error_pct"]) < 50, f"Density error > 50% for {model_name}" + + +@pytest.mark.parametrize("model_name", MODELS) +def test_compute_rdf(model_name: str) -> None: + """ + Compute and save Li-O RDFs and coordination numbers for one model. + + Parameters + ---------- + model_name : str + Name of the MLIP model in the registry. + """ + traj_path = find_trajectory(model_name) + if traj_path is None: + pytest.skip( + f"No NVT trajectory for {model_name}. Fetch the reference data or " + "generate it with test_reference_md (-m very_slow)." + ) + + first_frame = ase_read(str(traj_path), index=0, format="extxyz") + o_water_idx, o_tfsi_idx = identify_o_types(first_frame) + + assert len(o_water_idx) > 0, f"No O_water found for {model_name}" + assert len(o_tfsi_idx) > 0, f"No O_TFSI found for {model_name}" + + result = compute_rdf(traj_path, o_water_idx, o_tfsi_idx) + + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_dir / "rdf.json", "w") as f: + json.dump( + {k: v for k, v in result.items() if not isinstance(v, list)}, f, indent=2 + ) + np.savez( + out_dir / "rdf.npz", + r=np.array(result["r"]), + gr_LiO_total=np.array(result["gr_LiO_total"]), + gr_LiO_water=np.array(result["gr_LiO_water"]), + gr_LiO_TFSI=np.array(result["gr_LiO_TFSI"]), + ) + + assert 2.0 < result["coord_LiO_total"] < 8.0, ( + f"Unexpected Li-O_total CN={result['coord_LiO_total']:.2f} for {model_name}" + ) + + +@pytest.mark.parametrize("model_name", MODELS) +def test_compute_xray_sq(model_name: str) -> None: + """ + Compute and save X-ray S(q) in Faber-Ziman convention for one model. + + Parameters + ---------- + model_name : str + Name of the MLIP model in the registry. + """ + traj_path = find_trajectory(model_name) + if traj_path is None: + pytest.skip( + f"No NVT trajectory for {model_name}. Fetch the reference data or " + "generate it with test_reference_md (-m very_slow)." + ) + + result = compute_sq_travis_style(traj_path) + result["model"] = model_name + result["traj_path"] = str(traj_path) + + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_dir / "xray_sq.json", "w") as f: + json.dump(result, f, indent=2) + np.savez( + out_dir / "xray_sq.npz", + q=np.array(result["q"]), + Sq=np.array(result["Sq"]), + ) + + assert len(result["q"]) > 10, f"Too few q-points for {model_name}"