Ion-structure literature library

Note

The electronic and ionic states have been recalculated with default AA controls. Be MD uses the same IS-QOZ potential as HNC and VMHNC. See Otter 0.4.1 numerical validation.

This benchmark compares Otter ion-structure results with literature curves for aluminium, beryllium, and carbon.

In the four-panel \(S_{ii}(k)\) figure, panel 1 uses Gill et al., Fig. 3 [Gill et al., 2015]; panels 2 and 3 use Clérouin et al., Fig. 1 [Clérouin et al., 2015]; and panel 4 uses Wünsch et al., Fig. 2 [Wünsch et al., 2009]. For Wünsch Be, ordinary HNC, Rosenfeld–Ashcroft VMHNC and LAMMPS MD reuse one IS-QOZ pair potential. The MD error band is twice the standard error across RDF blocks or saved-frame reciprocal-shell averages. The real-space comparison uses Fig. 1(c). The carbon PA-HNC data were provided by C. E. Starrett (private communication; unpublished). Reference coordinates remain separate data files because they are digitized publication data; all unit conversions are explicit below.

The average-atom/pseudoatom construction follows Starrett and Saumon [2014], Starrett and Saumon [2013], and the default finite-temperature jellium LFC follows Chabrier [1990]. See the provenance, coordinate units, and data notice for the scientific interpretation. Both benchmark figures are exported as matching PNG and vector PDF files under benchmarks/outputs/ion_structure_library/figures.

Reproduction

From the root of the complete Otter checkout, using Poetry, run:

poetry run python benchmarks/examples/plot_ion_structure_library.py

Downloads are optional: .ipynb launches this repository script; .zip contains both formats. See Reproducing examples and benchmarks for setup.

The script calculates the states from their input parameters and then plots the results. No bundled Otter NPZ is required. Numerical outputs are written locally; literature reference tables remain inputs to the comparison.

Reproducing the MD curves also requires LAMMPS and MPI and can take hours. AA calculations keep one continuum worker per atom and use the default radial resolution and SCF tolerances.

Recorded results

The figures and output below are from the recorded validation run; running the source recalculates them with the installed Otter version.

Recorded result for plot_ion_structure_library Recorded result for plot_ion_structure_library

Recorded terminal output

state                                      model      obs reference                 RMSE        MAE        max
al_gill_rho2p7_te5_ti5                     Otter KS   sii Gill KS-PAMD        1.1147e-02 8.3394e-03 3.2241e-02
al_gill_rho2p7_te5_ti5                     Otter KS   sii Gill TF-PAMD        1.6657e-02 1.2596e-02 3.8828e-02
al_gill_rho2p7_te5_ti5                     Otter KS   sii Gill TF-DFT-MD      4.3173e-02 3.2838e-02 7.9956e-02
al_clerouin_rho8p1_te10_ti10               Otter KS   sii Clérouin OFMD       3.8447e-02 2.8862e-02 7.4195e-02
al_clerouin_rho8p1_te10_ti10               Otter TF   sii Clérouin OFMD       1.0938e-02 8.1710e-03 3.0392e-02
al_clerouin_rho8p1_te10_ti2                Otter KS   sii Clérouin OFMD       2.4673e-02 1.4629e-02 8.6048e-02
al_clerouin_rho8p1_te10_ti2                Otter KS   sii HNC-Y-SRR           4.2265e-02 3.2455e-02 1.1724e-01
al_clerouin_rho8p1_te10_ti2                Otter TF   sii Clérouin OFMD       1.9277e-02 1.1336e-02 8.0555e-02
al_clerouin_rho8p1_te10_ti2                Otter TF   sii HNC-Y-SRR           3.1915e-02 2.4906e-02 7.2462e-02
be_wunsch_rho5p544_te13_ti13               Otter-HNC  sii Wünsch DFT-MD       2.6001e-02 2.0192e-02 5.2515e-02
be_wunsch_rho5p544_te13_ti13               Otter-HNC  sii HNC-Y-SRR           1.6634e-02 1.3539e-02 3.2527e-02
be_wunsch_rho5p544_te13_ti13               Otter-HNC  sii HNC-KK              7.3300e-02 4.5546e-02 1.7372e-01
be_wunsch_rho5p544_te13_ti13               Otter-HNC  sii HNC-Y               7.4009e-02 5.4735e-02 1.3678e-01
be_wunsch_rho5p544_te13_ti13               Otter-HNC  gii Wünsch DFT-MD       6.3582e-02 4.6904e-02 1.3443e-01
be_wunsch_rho5p544_te13_ti13               Otter-HNC  gii HNC-Y-SRR           5.1892e-02 3.4911e-02 1.2601e-01
be_wunsch_rho5p544_te13_ti13               Otter-HNC  gii HNC-KK              4.5877e-02 3.3586e-02 8.8297e-02
be_wunsch_rho5p544_te13_ti13               Otter-HNC  gii HNC-Y               9.4556e-02 5.9866e-02 2.0475e-01
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC sii Wünsch DFT-MD       2.4126e-02 1.9148e-02 4.6549e-02
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC sii HNC-Y-SRR           1.4498e-02 1.1922e-02 2.8626e-02
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC sii HNC-KK              7.2075e-02 4.5361e-02 1.7164e-01
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC sii HNC-Y               7.7196e-02 5.7670e-02 1.4096e-01
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC gii Wünsch DFT-MD       5.9346e-02 4.4100e-02 1.2015e-01
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC gii HNC-Y-SRR           4.5971e-02 3.1079e-02 1.1087e-01
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC gii HNC-KK              4.2388e-02 3.0793e-02 8.3608e-02
be_wunsch_rho5p544_te13_ti13               Otter-VMHNC gii HNC-Y               1.0123e-01 6.3540e-02 2.2059e-01
be_wunsch_rho5p544_te13_ti13               Otter-MD   sii Wünsch DFT-MD       2.6313e-02 2.0924e-02 4.8405e-02
be_wunsch_rho5p544_te13_ti13               Otter-MD   sii HNC-Y-SRR           1.5119e-02 1.2679e-02 3.0992e-02
be_wunsch_rho5p544_te13_ti13               Otter-MD   sii HNC-KK              5.6188e-02 3.5164e-02 1.5141e-01
be_wunsch_rho5p544_te13_ti13               Otter-MD   sii HNC-Y               7.8817e-02 5.7852e-02 1.5011e-01
be_wunsch_rho5p544_te13_ti13               Otter-MD   gii Wünsch DFT-MD       6.0968e-02 4.4868e-02 1.1995e-01
be_wunsch_rho5p544_te13_ti13               Otter-MD   gii HNC-Y-SRR           4.6431e-02 3.1815e-02 1.0817e-01
be_wunsch_rho5p544_te13_ti13               Otter-MD   gii HNC-KK              3.8051e-02 2.8202e-02 7.4430e-02
be_wunsch_rho5p544_te13_ti13               Otter-MD   gii HNC-Y               1.0463e-01 6.6776e-02 2.2679e-01
c_starrett_rho20_te50_ti50                 Otter KS   gii Starrett PA-HNC     6.6633e-03 5.1869e-03 1.6273e-02
from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor, as_completed
from contextlib import ExitStack
import hashlib
import json
import os
from pathlib import Path
import sys
import time
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
from scipy.constants import physical_constants

from otter import (
    PlasmaWorkflowConfig,
    continue_plasma_workflow_from_electronic_result,
    solve_plasma_workflow,
)
from otter.numerics.constants import BOHR_TO_ANGSTROM
from otter.plotting import (
    MODEL_STYLES,
    grid_figsize,
    save_figure,
    style_context,
)


# =============================================================================
# User input
# =============================================================================
USE_PRECOMPUTED_DATA = False
if os.environ.get("OTTER_RECOMPUTE_ION_STRUCTURE_LIBRARY", "0") == "1":
    USE_PRECOMPUTED_DATA = False

# Three independent state groups, each with the default AA worker. The two
# Al 8.1-g/cc states share one electronic calculation because only the ion
# temperature differs.
MAX_STATE_WORKERS = 1
HNC_CLOSURE_TOL = 2.5e-3
R_RETAIN_MAX_BOHR = 20.0
K_RETAIN_MAX_BOHR_INV = 20.0
VMHNC_ETA_TOL = 1.0e-6
RUN_WUNSCH_SAME_POTENTIAL_MD = True
WUNSCH_MD_MPI_PROCESSES = 10
WUNSCH_MD_ATOMS = 2048
WUNSCH_MD_K_MAX_ANGSTROM_INV = 10.2
# The fundamental cubic-cell shell has only three independent half-space
# vectors and is visibly direction-starved.  Keep it in the NPZ audit data,
# but require at least the next shell's four modes when drawing MD S(k).
MD_MIN_HALF_SPACE_MODES_PER_BIN = 4
# =============================================================================



STATE_GROUPS: dict[str, tuple[dict[str, Any], ...]] = {
    "al_gill": (
        {
            "state_id": "al_gill_rho2p7_te5_ti5",
            "element": "Al",
            "rho_g_cc": 2.7,
            "te_ev": 5.0,
            "ti_ev": 5.0,
        },
    ),
    "al_clerouin": (
        {
            "state_id": "al_clerouin_rho8p1_te10_ti10",
            "element": "Al",
            "rho_g_cc": 8.1,
            "te_ev": 10.0,
            "ti_ev": 10.0,
        },
        {
            "state_id": "al_clerouin_rho8p1_te10_ti2",
            "element": "Al",
            "rho_g_cc": 8.1,
            "te_ev": 10.0,
            "ti_ev": 2.0,
        },
    ),
    "al_clerouin_tf": (
        {
            "state_id": "al_clerouin_rho8p1_te10_ti10_tf",
            "element": "Al",
            "rho_g_cc": 8.1,
            "te_ev": 10.0,
            "ti_ev": 10.0,
            "electronic_model": "tf",
        },
        {
            "state_id": "al_clerouin_rho8p1_te10_ti2_tf",
            "element": "Al",
            "rho_g_cc": 8.1,
            "te_ev": 10.0,
            "ti_ev": 2.0,
            "electronic_model": "tf",
        },
    ),
    "be_wunsch": (
        {
            "state_id": "be_wunsch_rho5p544_te13_ti13",
            "element": "Be",
            "rho_g_cc": 5.544,
            "te_ev": 13.0,
            "ti_ev": 13.0,
        },
    ),
    "c_starrett_hot": (
        {
            "state_id": "c_starrett_rho20_te50_ti50",
            "element": "C",
            "rho_g_cc": 20.0,
            "te_ev": 50.0,
            "ti_ev": 50.0,
        },
    ),
}

REFERENCE_SERIES: dict[str, tuple[dict[str, str], ...]] = {
    "al_gill_rho2p7_te5_ti5": (
        {
            "observable": "sii",
            "label": "Gill KS-PAMD",
            "file": "gill_et_al_2015/Sii_Al_T5ev_rho2.7_KS-PAMD_Gill.csv",
            "x_unit": "angstrom^-1",
        },
        {
            "observable": "sii",
            "label": "Gill TF-PAMD",
            "file": "gill_et_al_2015/Sii_Al_T5ev_rho2.7_TF-PAMD_Gill.csv",
            "x_unit": "angstrom^-1",
        },
        {
            "observable": "sii",
            "label": "Gill TF-DFT-MD",
            "file": "gill_et_al_2015/Sii_Al_T5ev_rho2.7_QMD_Gill.csv",
            "x_unit": "angstrom^-1",
        },
    ),
    "al_clerouin_rho8p1_te10_ti10": (
        {
            "observable": "sii",
            "label": "Clérouin OFMD",
            "file": ("clerouin_et_al_2015/" "Jean_2015_Al_rho8.1_Te10.0_Ti10.csv"),
            "x_unit": "angstrom^-1",
        },
    ),
    "al_clerouin_rho8p1_te10_ti2": (
        {
            "observable": "sii",
            "label": "Clérouin OFMD",
            "file": (
                "clerouin_et_al_2015/" "Jean_2015_Al_rho8.1_Te10.0_Ti2.0_OFMD.csv"
            ),
            "x_unit": "angstrom^-1",
        },
        {
            "observable": "sii",
            "label": "HNC-Y-SRR",
            "file": ("clerouin_et_al_2015/" "Jean_2015_Al_rho8.1_Te10.0_Ti2.0_SRR.csv"),
            "x_unit": "angstrom^-1",
        },
    ),
    "be_wunsch_rho5p544_te13_ti13": (
        {
            "observable": "sii",
            "label": "Wünsch DFT-MD",
            "file": ("wunsch_et_al_2009/" "Sii_DFTMD_Be_3rho0_T13_Z2_wunsch2009.csv"),
            "x_unit": "angstrom^-1",
        },
        {
            "observable": "sii",
            "label": "HNC-Y-SRR",
            "file": (
                "wunsch_et_al_2009/" "Sii_HNC-YSRR_Be_3rho0_T13_Z2_wunsch2009.csv"
            ),
            "x_unit": "angstrom^-1",
        },
        {
            "observable": "sii",
            "label": "HNC-KK",
            "file": ("wunsch_et_al_2009/" "Sii_HNCKK_Be_3rho0_T13_Z2_wunsch2009.csv"),
            "x_unit": "angstrom^-1",
        },
        {
            "observable": "sii",
            "label": "HNC-Y",
            "file": ("wunsch_et_al_2009/" "Sii_HNCY_Be_3rho0_T13_Z2_wunsch2009.csv"),
            "x_unit": "angstrom^-1",
        },
        {
            "observable": "gii",
            "label": "Wünsch DFT-MD",
            "file": ("wunsch_et_al_2009/" "gii_DFTMD_Be_3rho0_T13_Z2_wunsch2009.csv"),
            "x_unit": "angstrom",
        },
        {
            "observable": "gii",
            "label": "HNC-Y-SRR",
            "file": (
                "wunsch_et_al_2009/" "gii_HNC-YSRR_Be_3rho0_T13_Z2_wunsch2009.csv"
            ),
            "x_unit": "angstrom",
        },
        {
            "observable": "gii",
            "label": "HNC-KK",
            "file": ("wunsch_et_al_2009/" "gii_HNC-KK_Be_3rho0_T13_Z2_wunsch2009.csv"),
            "x_unit": "angstrom",
        },
        {
            "observable": "gii",
            "label": "HNC-Y",
            "file": ("wunsch_et_al_2009/" "gii_HNC-Y_Be_3rho0_T13_Z2_wunsch2009.csv"),
            "x_unit": "angstrom",
        },
    ),
    "c_starrett_rho20_te50_ti50": (
        {
            "observable": "gii",
            "label": "Starrett PA-HNC",
            "file": ("starrett_saumon_2013/" "gii_C_20gcc_50.0ev_starrett.csv"),
            "x_unit": "bohr",
        },
    ),
}

STATE_TITLES = {
    "al_gill_rho2p7_te5_ti5": (r"Al: $\rho=2.7$ g cm$^{-3}$, $T_e=T_i=5$ eV"),
    "al_clerouin_rho8p1_te10_ti10": (r"Al: $\rho=8.1$ g cm$^{-3}$, $T_e=T_i=10$ eV"),
    "al_clerouin_rho8p1_te10_ti2": (
        r"Al: $\rho=8.1$ g cm$^{-3}$, $T_e=10$, $T_i=2$ eV"
    ),
    "be_wunsch_rho5p544_te13_ti13": (r"Be: $\rho=5.544$ g cm$^{-3}$, $T_e=T_i=13$ eV"),
    "c_starrett_rho20_te50_ti50": (r"C: $\rho=20$ g cm$^{-3}$, $T_e=T_i=50$ eV"),
}

OTTER_SERIES = {
    state_id: ((state_id, "Otter KS", "-", ""),) for state_id in REFERENCE_SERIES
}
OTTER_SERIES.update(
    {
        "al_clerouin_rho8p1_te10_ti10": (
            ("al_clerouin_rho8p1_te10_ti10", "Otter KS", "-", ""),
            ("al_clerouin_rho8p1_te10_ti10_tf", "Otter TF", "--", ""),
        ),
        "al_clerouin_rho8p1_te10_ti2": (
            ("al_clerouin_rho8p1_te10_ti2", "Otter KS", "-", ""),
            ("al_clerouin_rho8p1_te10_ti2_tf", "Otter TF", "--", ""),
        ),
        "be_wunsch_rho5p544_te13_ti13": (
            ("be_wunsch_rho5p544_te13_ti13", "Otter-HNC", "-", ""),
            (
                "be_wunsch_rho5p544_te13_ti13",
                "Otter-VMHNC",
                "--",
                "vmhnc_",
            ),
            ("be_wunsch_rho5p544_te13_ti13", "Otter-MD", "-.", "md_"),
        ),
    }
)


def repository_root() -> Path:
    """Locate source and reference inputs, independently of numerical outputs."""
    candidates = [Path.cwd().resolve(), *Path.cwd().resolve().parents]
    source_file = globals().get("__file__")
    if source_file is not None:
        candidates.extend(Path(source_file).resolve().parents)
    for candidate in candidates:
        if (candidate / "pyproject.toml").is_file() and (candidate / "src/otter").is_dir():
            return candidate
    raise FileNotFoundError("Run from an Otter source checkout with its dependencies installed.")


ROOT = repository_root()
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))
PRECOMPUTED_DIR = ROOT / "benchmarks" / "baselines" / "ion_structure_library"
REFERENCE_DIR = ROOT / "benchmarks" / "reference_data" / "ion_structure_library"
OUTPUT_DIR = (
    ROOT / "benchmarks" / "outputs" / "ion_structure_library" / "gallery_recomputed"
)
FIGURE_DIR = ROOT / "benchmarks" / "outputs" / "ion_structure_library" / "figures"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def load_npz(path: Path) -> dict[str, np.ndarray]:
    with np.load(path, allow_pickle=False) as archive:
        state = {key: np.asarray(archive[key]) for key in archive.files}
    if any(value.dtype.hasobject for value in state.values()):
        raise TypeError(f"Object arrays are forbidden in {path}.")
    return state


def load_precomputed_states() -> dict[str, dict[str, np.ndarray]]:
    """Verify every accepted result checksum before plotting."""
    manifest = json.loads(
        (PRECOMPUTED_DIR / "manifest.json").read_text(encoding="utf-8")
    )
    if manifest.get("benchmark_id") != "ion_structure_library":
        raise ValueError("Unexpected ion-structure manifest.")
    loaded: dict[str, dict[str, np.ndarray]] = {}
    for record in manifest["states"]:
        path = PRECOMPUTED_DIR / str(record["baseline_file"])
        if sha256_file(path) != str(record["baseline_sha256"]):
            raise RuntimeError(f"Checksum mismatch for {path}.")
        state = load_npz(path)
        state_id = str(state["state_id"].item())
        if state_id != str(record["state_id"]):
            raise ValueError(f"State identifier mismatch in {path}.")
        loaded[state_id] = state
    return loaded


def workflow_config(
    state: dict[str, Any],
    *,
    ion_temperature_ev: float | None,
    bridge_model: str = "none",
) -> PlasmaWorkflowConfig:
    """Build the complete public Otter workflow for one thermodynamic state."""
    model = str(state.get("electronic_model", "qm"))
    model_override = {} if model == "qm" else {"electronic_model": model}
    return PlasmaWorkflowConfig(
        elements=[str(state["element"])],
        temperature_ev=float(state["te_ev"]),
        ion_temperature_ev=(
            None if ion_temperature_ev is None else float(ion_temperature_ev)
        ),
        rho_g_cc=float(state["rho_g_cc"]),
        **model_override,
        hnc_closure_transform_tol=float(HNC_CLOSURE_TOL),
        hnc_bridge_model=bridge_model,
        **({"vmhnc_eta_tol": VMHNC_ETA_TOL} if bridge_model != "none" else {}),
    )


def strict_check(
    workflow: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Reject unconverged electronic or HNC best-effort output."""
    electronic = dict(workflow["electronic"]["result"])
    ion = dict(workflow["ion"])
    from benchmarks.runners.ion_structure_validation import validate_screening_charge
    validate_screening_charge(ion)
    if electronic.get("stage2_converged") is not True:
        raise RuntimeError("Full average-atom stage 2 did not converge.")
    if dict(electronic.get("ext_status", {})).get("converged") is not True:
        raise RuntimeError("External fixed-mu average atom did not converge.")
    if str(electronic.get("threshold_state_status", "")).lower() == "unresolved":
        raise RuntimeError("The threshold-state representation is unresolved.")
    if ion.get("hnc_converged") is not True:
        raise RuntimeError("HNC did not reach a physical fixed point.")
    if float(ion["hnc_output_residual"]) > float(
        workflow["configuration"]["hnc_tol"]
    ):
        raise RuntimeError("HNC residual exceeds the configured tolerance.")
    if float(ion["closure_transform_max_abs"]) > HNC_CLOSURE_TOL:
        raise RuntimeError("The g/S transform-closure audit failed.")
    return electronic, ion


def pack_result(
    workflow: dict[str, Any],
    state: dict[str, Any],
    *,
    elapsed_s: float,
) -> dict[str, np.ndarray]:
    """Keep the portable arrays needed by this benchmark and downstream use."""
    electronic, ion = strict_check(workflow)
    r = np.asarray(ion["r"], dtype=float)
    k = np.asarray(ion["k"], dtype=float)
    r_mask = r <= R_RETAIN_MAX_BOHR
    k_mask = k <= K_RETAIN_MAX_BOHR_INV
    return {
        "schema_version": np.asarray("otter_gallery_ion_library_v2"),
        "storage_profile": np.asarray("benchmark_analysis"),
        "state_id": np.asarray(str(state["state_id"])),
        "electronic_model": np.asarray(str(state.get("electronic_model", "qm"))),
        "element": np.asarray(str(state["element"])),
        "rho_g_cc": np.asarray(float(state["rho_g_cc"])),
        "te_ev": np.asarray(float(state["te_ev"])),
        "ti_ev": np.asarray(float(state["ti_ev"])),
        "producer_elapsed_s": np.asarray(float(elapsed_s)),
        "n0_bohr3": np.asarray(float(electronic["n0"])),
        "r_ws_bohr": np.asarray(float(electronic["r_ws"])),
        "mu_ha": np.asarray(float(electronic["mu"])),
        "r_bohr": r[r_mask],
        "gii_r": np.asarray(ion["gij_r"][0, 0])[r_mask],
        "k_bohr_inv": k[k_mask],
        "sii_k": np.asarray(ion["sij_k"][0, 0])[k_mask],
        "zbar_partition": np.asarray(np.asarray(ion["zbar_partition"]).item()),
        "hnc_best_residual": np.asarray(float(ion["hnc_output_residual"])),
        "hnc_closure_mismatch": np.asarray(float(ion["closure_transform_max_abs"])),
    }


def add_wunsch_vmhnc(
    payload: dict[str, np.ndarray],
    hnc_workflow: dict[str, Any],
    vmhnc_workflow: dict[str, Any],
) -> None:
    """Store the second closure after checking that its potential is unchanged."""
    _, hnc = strict_check(hnc_workflow)
    _, vmhnc = strict_check(vmhnc_workflow)
    if str(hnc["hnc_bridge_model"]) != "none":
        raise RuntimeError("The ordinary-HNC Wünsch curve used a bridge.")
    if str(vmhnc["hnc_bridge_model"]) != "rosenfeld_ashcroft":
        raise RuntimeError("The Wünsch VMHNC curve used the wrong bridge.")
    if not np.allclose(hnc["vii_r"], vmhnc["vii_r"], rtol=0.0, atol=0.0):
        raise RuntimeError("Wünsch HNC and VMHNC did not reuse one potential.")
    r = np.asarray(vmhnc["r"], dtype=float)
    k = np.asarray(vmhnc["k"], dtype=float)
    r_mask, k_mask = r <= R_RETAIN_MAX_BOHR, k <= K_RETAIN_MAX_BOHR_INV
    payload.update(
        {
            "vmhnc_gii_r": np.asarray(vmhnc["gij_r"][0, 0])[r_mask],
            "vmhnc_sii_k": np.asarray(vmhnc["sij_k"][0, 0])[k_mask],
            "vmhnc_best_residual": np.asarray(float(vmhnc["hnc_output_residual"])),
            "vmhnc_closure_mismatch": np.asarray(
                float(vmhnc["closure_transform_max_abs"])
            ),
            "vmhnc_eta": np.asarray(float(vmhnc["vmhnc_eta"])),
            "vmhnc_variational_residual": np.asarray(
                float(vmhnc["vmhnc_variational_residual"])
            ),
        }
    )


def add_wunsch_md(
    payload: dict[str, np.ndarray],
    workflow: dict[str, Any],
) -> None:
    """Run the reusable single-/multi-species LAMMPS driver for Wünsch Be."""
    from tools.otter_lammps_md import (
        MDConfig,
        MDSpecies,
        PairPotential,
        run_otter_lammps_md,
    )

    ion = dict(workflow["ion"])
    mass_u = 9.0121831
    mass_ratio = (
        physical_constants["atomic mass constant"][0]
        / physical_constants["electron mass"][0]
    )
    atomic_time_ps = physical_constants["atomic unit of time"][0] * 1.0e12
    omega_p = np.sqrt(
        4.0
        * np.pi
        * np.asarray(ion["n_i"]).item()
        * np.asarray(ion["zbar_partition"]).item() ** 2
        / (mass_u * mass_ratio)
    )
    dt_omega = max(0.005, 0.01e-3 / atomic_time_ps * omega_p)
    production_steps = round(500.0 / dt_omega)
    state_id = str(payload["state_id"].item())
    result = run_otter_lammps_md(
        MDConfig(
            output_dir=OUTPUT_DIR / "md_work" / state_id,
            species=(MDSpecies("Be", mass_u, WUNSCH_MD_ATOMS),),
            ion_density_bohr3=np.asarray(ion["n_i"]).item(),
            ion_temperature_ev=float(payload["ti_ev"]),
            timestep_ps=dt_omega / omega_p * atomic_time_ps,
            thermostat_damp_ps=0.5 / omega_p * atomic_time_ps,
            equilibration_steps=round(50.0 / dt_omega),
            production_steps=production_steps,
            rdf_bins=500,
            rdf_every=max(1, production_steps // 1000),
            rdf_repeat=50,
            trajectory_every=max(1, production_steps // 20),
            table_points=8_192,
            k_max_angstrom_inv=WUNSCH_MD_K_MAX_ANGSTROM_INV,
            random_seed=20_260_825,
            mpi_processes=WUNSCH_MD_MPI_PROCESSES,
            structure_factor_workers=WUNSCH_MD_MPI_PROCESSES,
        ),
        [PairPotential("Be", "Be", np.asarray(ion["r"]), np.asarray(ion["vii_r"]))],
    )
    payload.update(result)
    payload["md_gii_r"] = np.asarray(result["md_gij_r"])[:, 0]
    payload["md_gii_block_sem"] = np.asarray(result["md_gij_block_sem"])[:, 0]
    payload["md_sii_k"] = np.asarray(result["md_snn_k"])
    payload["md_sii_frame_sem"] = np.asarray(result["md_snn_frame_sem"])
    payload["md_sii_vectors_per_bin"] = np.asarray(result["md_vectors_per_k_bin"])
    for key in (
        "md_gij_r",
        "md_gij_block_sem",
        "md_snn_k",
        "md_snn_frame_sem",
        "md_sij_k",
        "md_sij_frame_sem",
        "md_vectors_per_k_bin",
    ):
        payload.pop(key, None)


def solve_group(
    group: tuple[dict[str, Any], ...],
) -> dict[str, dict[str, np.ndarray]]:
    """Solve one electronic state and reuse it when only Ti changes."""
    first = group[0]
    started = time.perf_counter()
    electronic_only = solve_plasma_workflow(
        workflow_config(first, ion_temperature_ev=None)
    )
    electronic_elapsed = time.perf_counter() - started
    electronic_kind = str(electronic_only["electronic"]["kind"])
    electronic = dict(electronic_only["electronic"]["result"])

    solved: dict[str, dict[str, np.ndarray]] = {}
    for state in group:
        ion_started = time.perf_counter()
        workflow = continue_plasma_workflow_from_electronic_result(
            workflow_config(
                state,
                ion_temperature_ev=float(state["ti_ev"]),
            ),
            electronic_kind=electronic_kind,
            electronic_result=electronic,
        )
        elapsed = electronic_elapsed + time.perf_counter() - ion_started
        payload = pack_result(workflow, state, elapsed_s=elapsed)
        if str(state["state_id"]).startswith("be_wunsch"):
            vmhnc = continue_plasma_workflow_from_electronic_result(
                workflow_config(
                    state,
                    ion_temperature_ev=float(state["ti_ev"]),
                    bridge_model="rosenfeld_ashcroft",
                ),
                electronic_kind=electronic_kind,
                electronic_result=electronic,
            )
            add_wunsch_vmhnc(payload, workflow, vmhnc)
            if RUN_WUNSCH_SAME_POTENTIAL_MD:
                add_wunsch_md(payload, workflow)
        solved[str(state["state_id"])] = payload
    return solved


def solve_all_states() -> dict[str, dict[str, np.ndarray]]:
    """Calculate all state groups with a bounded process pool."""
    groups = tuple(STATE_GROUPS.values())
    loaded: dict[str, dict[str, np.ndarray]] = {}
    with ProcessPoolExecutor(max_workers=MAX_STATE_WORKERS) as pool:
        futures = {pool.submit(solve_group, group): group for group in groups}
        for future in as_completed(futures):
            group_result = future.result()
            loaded.update(group_result)
            print("[computed] " + ", ".join(sorted(group_result)))

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    for state_id, state in loaded.items():
        path = OUTPUT_DIR / f"{state_id}.npz"
        np.savez_compressed(path, **state)
        print(f"[saved] {path}")
    return loaded


def load_reference(series: dict[str, str]) -> tuple[np.ndarray, np.ndarray]:
    """Load one two-column digitization without changing its coordinate unit."""
    path = (REFERENCE_DIR / series["file"]).resolve()
    if not path.is_relative_to(REFERENCE_DIR.resolve()):
        raise ValueError("Reference path escapes its package directory.")
    values = np.asarray(
        np.genfromtxt(path, delimiter=",", comments="#"),
        dtype=float,
    )
    if values.ndim != 2 or values.shape[1] < 2:
        raise ValueError(f"Expected two reference columns in {path}.")
    mask = np.isfinite(values[:, 0]) & np.isfinite(values[:, 1])
    return values[mask, 0], values[mask, 1]


def otter_curve(
    state: dict[str, np.ndarray],
    observable: str,
    x_unit: str,
    prefix: str = "",
) -> tuple[np.ndarray, np.ndarray]:
    """Convert only the Otter coordinate to the publication's stated unit."""
    coordinate_prefix = "" if prefix == "vmhnc_" else prefix
    if observable == "sii":
        x = np.asarray(state[f"{coordinate_prefix}k_bohr_inv"], dtype=float)
        y = np.asarray(state[f"{prefix}sii_k"], dtype=float)
        if prefix == "md_":
            reliable = (
                np.asarray(state["md_sii_vectors_per_bin"], dtype=int)
                >= MD_MIN_HALF_SPACE_MODES_PER_BIN
            )
            x, y = x[reliable], y[reliable]
        if x_unit == "angstrom^-1":
            x = x / BOHR_TO_ANGSTROM
        elif x_unit != "bohr^-1":
            raise ValueError(f"Unsupported reciprocal unit {x_unit!r}.")
        return x, y
    if observable == "gii":
        x = np.asarray(state[f"{coordinate_prefix}r_bohr"], dtype=float)
        y = np.asarray(state[f"{prefix}gii_r"], dtype=float)
        if x_unit == "angstrom":
            x = x * BOHR_TO_ANGSTROM
        elif x_unit != "bohr":
            raise ValueError(f"Unsupported radius unit {x_unit!r}.")
        return x, y
    raise ValueError(f"Unknown observable {observable!r}.")


def print_metrics(states: dict[str, dict[str, np.ndarray]]) -> None:
    if any(bool(state.get("md_is_historical", False)) for state in states.values()):
        print("MD rows are historical old-potential results; Otter HNC/VMHNC curves were refreshed.")
    print(
        f"{'state':42s} {'model':10s} {'obs':>3s} {'reference':19s} "
        f"{'RMSE':>10s} {'MAE':>10s} {'max':>10s}"
    )
    for state_id, series_list in REFERENCE_SERIES.items():
        for result_id, model_label, _, prefix in OTTER_SERIES[state_id]:
            for series in series_list:
                x_ref, y_ref = load_reference(series)
                x_otter, y_otter = otter_curve(
                    states[result_id],
                    str(series["observable"]),
                    str(series["x_unit"]),
                    prefix,
                )
                mask = (x_ref >= x_otter[0]) & (x_ref <= x_otter[-1])
                delta = np.interp(x_ref[mask], x_otter, y_otter) - y_ref[mask]
                print(
                    f"{state_id:42s} {model_label:10s} "
                    f"{series['observable']:>3s} "
                    f"{series['label'][:19]:19s} "
                    f"{np.sqrt(np.mean(delta**2)):10.4e} "
                    f"{np.mean(np.abs(delta)):10.4e} "
                    f"{np.max(np.abs(delta)):10.4e}"
                )


def main() -> None:
    states = load_precomputed_states() if USE_PRECOMPUTED_DATA else solve_all_states()
    print(
        "Using "
        + (
            "checksummed, precomputed Otter results."
            if USE_PRECOMPUTED_DATA
            else "new results calculated directly by this gallery script."
        )
    )
    print_metrics(states)


    def plot_observable(
        *,
        observable: str,
        state_ids: tuple[str, ...],
    ) -> plt.Figure:
        """Draw one self-contained comparison figure."""
        ncols = 2
        nrows = (len(state_ids) + ncols - 1) // ncols
        fig, axes = plt.subplots(
            nrows,
            ncols,
            figsize=grid_figsize(nrows, ncols),
            squeeze=False,
        )
        marker_cycle = ("o", "s", "^", "x")
        colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
        for panel, state_id in enumerate(state_ids):
            axis = axes.ravel()[panel]
            series_list = [
                item
                for item in REFERENCE_SERIES[state_id]
                if item["observable"] == observable
            ]
            display_unit = str(series_list[0]["x_unit"])
            for model_index, (result_id, label, line_style, prefix) in enumerate(
                OTTER_SERIES[state_id]
            ):
                x_otter, y_otter = otter_curve(
                    states[result_id],
                    observable,
                    display_unit,
                    prefix,
                )
                style = dict(MODEL_STYLES["otter"])
                style["linestyle"] = line_style
                style["alpha"] = 0.82
                if model_index == 1:
                    style["color"] = colors[0]
                elif model_index == 2:
                    style["color"] = colors[1]
                if prefix == "md_":
                    style.update(linestyle="--", marker="o", markersize=2.5,
                                 markevery=10)
                if prefix == "md_" and bool(states[result_id].get("md_is_historical", False)):
                    label += " (old potential)"
                axis.plot(x_otter, y_otter, label=label, **style)
                sem_key = (
                    "md_sii_frame_sem"
                    if prefix == "md_" and observable == "sii"
                    else "md_gii_block_sem"
                )
                if prefix == "md_" and sem_key in states[result_id]:
                    sem = np.asarray(states[result_id][sem_key], dtype=float)
                    if observable == "sii":
                        reliable = (
                            np.asarray(
                                states[result_id]["md_sii_vectors_per_bin"],
                                dtype=int,
                            )
                            >= MD_MIN_HALF_SPACE_MODES_PER_BIN
                        )
                        sem = sem[reliable]
                    axis.fill_between(
                        x_otter,
                        y_otter - 2.0 * sem,
                        y_otter + 2.0 * sem,
                        color=style["color"],
                        alpha=0.14,
                        linewidth=0.0,
                    )
            reference_x: list[np.ndarray] = []
            for index, series in enumerate(series_list):
                x_ref, y_ref = load_reference(series)
                reference_x.append(x_ref)
                marker = marker_cycle[index % len(marker_cycle)]
                scatter_options: dict[str, Any] = {
                    "s": 25,
                    "marker": marker,
                    "linewidths": 1.2,
                    "label": series["label"],
                    "zorder": 3,
                }
                color = colors[index % len(colors)]
                if marker == "x":
                    scatter_options["color"] = color
                else:
                    scatter_options["facecolors"] = "none"
                    scatter_options["edgecolors"] = color
                axis.scatter(x_ref, y_ref, **scatter_options)
            axis.set_title(STATE_TITLES[state_id], fontsize=10)
            axis.set_ylabel(r"$S_{ii}(k)$" if observable == "sii" else r"$g_{ii}(r)$")
            if display_unit == "angstrom^-1":
                axis.set_xlabel(r"$k$ [$\mathrm{\AA}^{-1}$]")
            elif display_unit == "angstrom":
                axis.set_xlabel(r"$r$ [$\mathrm{\AA}$]")
            else:
                axis.set_xlabel(r"$r$ [Bohr]")
            all_reference_x = np.concatenate(reference_x)
            span = float(np.ptp(all_reference_x))
            margin = max(0.03 * span, 1.0e-6)
            left = (
                -0.5
                if observable == "gii"
                else max(0.0, float(np.min(all_reference_x)) - margin)
            )
            axis.set_xlim(left, float(np.max(all_reference_x)) + margin)
            axis.axhline(1.0, color="0.55", lw=0.8, ls=":")
            axis.legend(fontsize="small")
        for panel in range(len(state_ids), axes.size):
            axes.ravel()[panel].set_visible(False)
        fig.suptitle("Otter ion structure versus curated literature curves", y=0.985)
        source_line = (
            "Reference data: Gill et al. (2015); Clérouin et al. (2015); "
            "Wunsch et al. (2009)."
            if observable == "sii"
            else "Reference data: Wunsch et al. (2009); "
            "C. E. Starrett (private communication)."
        )
        fig.text(
            0.5,
            0.006,
            source_line,
            ha="center",
            va="bottom",
            fontsize=7.5,
        )
        fig.tight_layout(rect=(0.0, 0.035, 1.0, 0.965), pad=0.55)
        return fig


    # %%
    # Static ion structure factors
    # ----------------------------
    #
    # Literature wave numbers are stored in inverse ångström.  Otter's native
    # inverse-Bohr grid is converted explicitly inside ``otter_curve``.

    plot_style = ExitStack()
    plot_style.enter_context(style_context("thesis", palette="bing"))
    fig_sii = plot_observable(
        observable="sii",
        state_ids=(
            "al_gill_rho2p7_te5_ti5",
            "al_clerouin_rho8p1_te10_ti10",
            "al_clerouin_rho8p1_te10_ti2",
            "be_wunsch_rho5p544_te13_ti13",
        ),
    )


    # %%
    # Pair distribution functions
    # ---------------------------
    #
    # The Be coordinates are ångström, while the carbon digitization uses Bohr.

    fig_gii = plot_observable(
        observable="gii",
        state_ids=(
            "be_wunsch_rho5p544_te13_ti13",
            "c_starrett_rho20_te50_ti50",
        ),
    )

    save_figure(
        fig_sii,
        FIGURE_DIR / "ion_structure_library_sii",
        close=False,
    )
    save_figure(
        fig_gii,
        FIGURE_DIR / "ion_structure_library_gii",
        close=False,
    )
    plot_style.close()

    plt.show()



if __name__ == "__main__":
    main()

Gallery generated by Sphinx-Gallery