"""J-value extraction from a completed four-state DFT campaign.
Pipeline (mirrors what the user did manually before this module existed):
1. **Parse** the ``four_state_report.dat`` written by
:func:`mag4.cli.fourstate.main` to recover, for each shell ``Jₙ``,
which 4 unique configurations contribute to the formula
``JS² = (E_uu + E_dd − E_ud − E_du) / 4``.
2. **Read** each configuration's converged DFT total energy from
``configN/`` via :func:`mag4.energy.get_energy`.
3. **Combine** energies with the formulas to produce ``JS²`` (and ``J``
if a spin S is given).
This module is the pure-Python core; the CLI in :mod:`mag4.cli.extract`
adds argument parsing, pretty-printing and ``j_values.dat`` output.
"""
from __future__ import annotations
import os
import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from mag4.constants import KB_MEV
from mag4.energy import EnergyResult, get_energy
#: Name of the report file produced by ``mag4-magnetic`` (default).
DEFAULT_REPORT_NAME = "four_state_report.dat"
# ---------------------------------------------------------------------------
# Report-file parsing
# ---------------------------------------------------------------------------
# Matches: " J1 (d=3.8987 Å) = ( E[#1] + E[#2] - E[#3] - E[#4] ) / 4"
# (the writer is in mag4.reports.print_report; format is fixed)
_FORMULA_RE = re.compile(
r"^\s*(?P<name>J\d+)\s*"
r"\(d=(?P<dist>[\d.+-]+)\s*Å\)\s*=\s*"
r"\(\s*"
r"E\[#(?P<uu>\d+)\]\s*\+\s*"
r"E\[#(?P<dd>\d+)\]\s*-\s*"
r"E\[#(?P<ud>\d+)\]\s*-\s*"
r"E\[#(?P<du>\d+)\]\s*"
r"\)\s*/\s*(?P<div>\d+)\s*$" # divisor = 4·multiplicity (4 if M=1)
)
# Matches: " J8 (d=5.5439 Å): INCOMPLETE -- only 0 states found. ..."
_INCOMPLETE_RE = re.compile(
r"^\s*(?P<name>J\d+)\s*\(d=(?P<dist>[\d.+-]+)\s*Å\)\s*:\s*INCOMPLETE",
)
# Matches: " Config #5: J3(dd), J5(dd)"
_CONFIG_HEADER_RE = re.compile(r"^\s*Config\s*#(?P<uid>\d+)\s*:")
# Matches: " Cell: a=4.195 | b=4.195 | c=8.390 | … | n_atoms=16 | …"
_CONFIG_CELL_RE = re.compile(r"n_atoms\s*=\s*(?P<n>\d+)")
# Matches: " Total atoms in supercell : 16"
_SUPERCELL_ATOMS_RE = re.compile(
r"Total atoms in supercell\s*:\s*(?P<n>\d+)", re.IGNORECASE
)
[docs]
def parse_report_config_cells(report_path: str) -> Dict[int, int]:
"""Read per-config ``n_atoms`` values from a four_state_report.dat.
The report emits a ``Cell: a=… | … | n_atoms=N | …`` line under each
``Config #uid:`` block when WIEN2k cell reduction is on. Returns
``{config_index → n_atoms}``. Older reports without the line yield
an empty dict and the caller falls back to assuming every config
sits in the same cell (today's behaviour, no rescaling).
"""
out: Dict[int, int] = {}
if not os.path.isfile(report_path):
return out
pending_uid: Optional[int] = None
with open(report_path, encoding="utf-8", errors="replace") as fh:
for line in fh:
m = _CONFIG_HEADER_RE.match(line)
if m:
pending_uid = int(m.group("uid"))
continue
if pending_uid is not None:
m = _CONFIG_CELL_RE.search(line)
if m:
out[pending_uid] = int(m.group("n"))
pending_uid = None
return out
[docs]
def read_wien2k_struct_atom_count(struct_path: str) -> Optional[int]:
"""Sum the ``MULT=N`` fields across a WIEN2k ``case.struct`` file.
This is the authoritative atom count WIEN2k's ``:ENE`` refers to.
Whenever ``init_lapw``'s ``sgroup`` rewrites the struct (cell
reduction, lattice-letter correction), mag4's pre-emption from
``four_state_report.dat`` can disagree -- the post-``sgroup``
file on disk is the truth. Returns ``None`` if the file is
missing or no ``MULT=`` field could be parsed.
Example: a P6_3/mmc primitive cell with 2 Cr + 2 Sb listed
(each MULT=2) returns 4.
"""
if not os.path.isfile(struct_path):
return None
total = 0
try:
with open(struct_path, encoding="utf-8", errors="replace") as fh:
for line in fh:
m = re.search(r"\bMULT=\s*(\d+)", line)
if m:
total += int(m.group(1))
except OSError: # pragma: no cover - read error
return None
return total if total > 0 else None
[docs]
def read_struct_atom_counts(compound_dir: str,
n_configs: int) -> Dict[int, int]:
"""For every ``configN/`` directory under ``compound_dir``, sum the
``MULT=`` fields of ``configN.struct`` (whatever WIEN2k's
``init_lapw`` left on disk after possibly accepting ``sgroup``'s
rewrite) and return ``{config_index → atom_count}``.
Missing struct files or unparseable ones are simply skipped --
callers should fall back on the report's recorded counts.
"""
out: Dict[int, int] = {}
for i in range(1, n_configs + 1):
cfg_dir = os.path.join(compound_dir, f"config{i}")
# Prefer ``configN.struct`` (= post-init / post-sgroup) over
# ``configN.struct_sgroup`` (the sgroup-rewritten file that
# job.init then copies onto ``configN.struct``).
for name in (f"config{i}.struct",
f"config{i}.struct_sgroup"):
n = read_wien2k_struct_atom_count(os.path.join(cfg_dir, name))
if n is not None:
out[i] = n
break
return out
[docs]
def parse_report_supercell_atoms(report_path: str) -> Optional[int]:
"""Read the total atom count in the supercell from a
four_state_report.dat header.
The four-state formula is defined on the supercell mag4 built
(e.g. NiO 1×1×2 → 16 atoms). When configs are computed in
*reduced* cells (Fm-3m conv → 2 atoms, etc.), their WIEN2k :ENE
values must be scaled **up to the supercell size** before
applying the formula. This helper returns the integer that
``mag4-extract`` should use as the scaling target.
Returns ``None`` for reports written by mag4 versions before
the ``Total atoms in supercell`` line was added; the caller
then falls back to ``math.lcm(per-config n_atoms)``.
"""
if not os.path.isfile(report_path):
return None
with open(report_path, encoding="utf-8", errors="replace") as fh:
for line in fh:
m = _SUPERCELL_ATOMS_RE.search(line)
if m:
return int(m.group("n"))
return None
_SUPERCELL_DIMS_RE = re.compile(
r"Supercell\s*:\s*(\d+)\s*[x×]\s*(\d+)\s*[x×]\s*(\d+)", re.IGNORECASE
)
[docs]
def parse_report_supercell(report_path: str) -> Optional[Tuple[int, int, int]]:
"""Read the supercell multipliers ``(na, nb, nc)`` from a four-state report
header line ``Supercell : na×nb×nc``. ``None`` if absent."""
if not os.path.isfile(report_path):
return None
with open(report_path, encoding="utf-8", errors="replace") as fh:
for line in fh:
m = _SUPERCELL_DIMS_RE.search(line)
if m:
return (int(m.group(1)), int(m.group(2)), int(m.group(3)))
return None
[docs]
def parse_report_config_sz(report_path: str) -> Dict[int, int]:
"""``{config_index: Sz}`` from the per-config ``Spins:`` lines, where
``Sz = Σ s_i`` of the ±1 spin vector (the net collinear moment quantum
number of that configuration)."""
out: Dict[int, int] = {}
if not os.path.isfile(report_path):
return out
cur: Optional[int] = None
with open(report_path, encoding="utf-8", errors="replace") as fh:
for line in fh:
m = _CONFIG_HEADER_RE.search(line)
if m:
cur = int(m.group("uid"))
continue
if cur is not None and "Spins:" in line:
spins = [int(x) for x in re.findall(r"[+-]?1",
line.split("Spins:", 1)[1])]
out[cur] = sum(spins)
cur = None
return out
_TR_PAIR_RE = re.compile(r"pair:\s*config(\d+)\s*==\s*config(\d+)")
[docs]
def parse_report_tr_pairs(report_path: str) -> List[Tuple[int, int]]:
"""``[(configA, configB), …]`` from the TIME-REVERSAL DEGENERACY CHECK
section written by ``mag4-magnetic --check-degeneracy`` (any method).
Each pair is related only by an antiunitary operation (unitary crystal op
× time reversal), so the two DFT energies must coincide; the caller
reports ``|ΔE|`` per pair. Empty when the report has no such section.
"""
pairs: List[Tuple[int, int]] = []
if not os.path.isfile(report_path):
return pairs
with open(report_path, encoding="utf-8", errors="replace") as fh:
for line in fh:
m = _TR_PAIR_RE.search(line)
if m:
pairs.append((int(m.group(1)), int(m.group(2))))
return pairs
# ---------------------------------------------------------------------------
# J-value computation
# ---------------------------------------------------------------------------
[docs]
@dataclass
class JResult:
"""One row of the final J table emitted by :func:`compute_j_values`.
Attributes
----------
name : str
Coupling label (``"J1"`` …).
distance_A : float
Mean M–M distance of the shell, Å.
formula : Formula
The original parsed formula (provides config indices, raw text).
energies_eV : dict[str, float or None]
Per-state total energies in eV. Keys: ``"uu"``, ``"dd"``,
``"ud"``, ``"du"``. ``None`` when the corresponding config did
not converge (or didn't exist).
JS2_eV, JS2_meV : float or None
Computed ``JS² = (E_uu + E_dd − E_ud − E_du) / 4``. ``None`` if
the formula could not be evaluated (any energy missing).
J_meV, J_K : float or None
Computed ``J = JS² / S²`` in meV and Kelvin (``J / k_B``).
Only populated when a positive spin ``S`` was passed.
error : str or None
Human-readable reason if ``JS2_eV`` is ``None`` (which configs
failed, INCOMPLETE shell, …).
"""
name: str
distance_A: float
formula: Formula
energies_eV: Dict[str, Optional[float]] = field(default_factory=dict)
JS2_eV: Optional[float] = None
JS2_meV: Optional[float] = None
J_meV: Optional[float] = None
J_K: Optional[float] = None
error: Optional[str] = None
@property
def computed(self) -> bool:
"""True iff ``JS² could be computed`` (all 4 energies present)."""
return self.JS2_eV is not None
[docs]
def compute_j_values(formulas: List[Formula],
energies: Dict[int, EnergyResult],
spin: Optional[float] = None,
config_n_atoms: Optional[Dict[int, int]] = None,
n_supercell_atoms: Optional[int] = None,
) -> List[JResult]:
"""Apply each formula to the matching energies and return the J table.
Parameters
----------
formulas : list of :class:`Formula`
Output of :func:`parse_report_formulas`.
energies : dict[int, EnergyResult]
Map ``config_index → EnergyResult`` for **every** config the
formulas reference. Indices come from the ``configN/`` dir
name, 1-based. Configs that didn't converge MUST still be
present with ``EnergyResult.converged = False`` so this
function can attribute the failure correctly.
spin : float, optional
Magnetic spin quantum number ``S``. When given (>0), each
``JS²`` is divided by ``S²`` to also report ``J`` in meV and
in Kelvin.
config_n_atoms : dict[int, int], optional
Map ``config_index → n_atoms`` from
:func:`parse_report_config_cells`. When supplied, each
contributing config's total energy is rescaled by
``N_LCM / N_X`` before applying the four-state formula —
necessary when ``mag4-magnetic`` wrote each config in its own
primitive cell (``--reduce-cell``, the default). When omitted
every config is assumed to sit in the same cell (legacy
behaviour, no rescaling).
Returns
-------
list of :class:`JResult`
One entry per input formula, same order.
"""
import math
results: List[JResult] = []
for f in formulas:
r = JResult(name=f.name, distance_A=f.distance_A, formula=f)
if f.incomplete:
r.error = "INCOMPLETE shell (cell too small — see four_state_report.dat)"
results.append(r)
continue
bad: List[str] = []
for label, idx in zip(("uu", "dd", "ud", "du"), f.config_indices):
er = energies.get(idx)
if er is None:
r.energies_eV[label] = None
bad.append(f"#{idx} (no EnergyResult provided)")
continue
if er.energy_eV is None or not er.converged:
r.energies_eV[label] = er.energy_eV
why = (er.diagnostics[0] if er.diagnostics
else "no energy / not converged")
bad.append(f"#{idx} ({why})")
continue
r.energies_eV[label] = er.energy_eV
if bad:
r.error = "missing/non-converged config(s): " + "; ".join(bad)
results.append(r)
continue
e = r.energies_eV # all four are populated and finite here
# Per-config cell scaling (DFT energy is extensive):
# E_X expressed in an N_target-atom cell is E_X · (N_target / N_X)
# Baselines cancel pairwise in the four-state formula → exact.
#
# ``N_target`` is preferentially the working supercell mag4 built
# (read from the report's ``Total atoms in supercell`` line) —
# that's the cell the four-state formula was *defined* against.
# If absent (old report) we fall back to ``math.lcm(N_X)``, which
# is numerically equivalent for J extraction but loses the link
# to the user-visible supercell size.
if config_n_atoms:
Ns = [config_n_atoms.get(idx) for idx in f.config_indices]
if all(n is not None and n > 0 for n in Ns):
if n_supercell_atoms is not None and n_supercell_atoms > 0:
n_target = n_supercell_atoms
else:
n_target = math.lcm(*Ns)
e = {
label: r.energies_eV[label] * (n_target / Ns[k])
for k, label in enumerate(("uu", "dd", "ud", "du"))
}
JS2_eV = (e["uu"] + e["dd"] - e["ud"] - e["du"]) / float(f.divisor) # type: ignore[operator]
r.JS2_eV = JS2_eV
r.JS2_meV = JS2_eV * 1000.0
if spin is not None and spin > 0:
r.J_meV = r.JS2_meV / (spin * spin)
r.J_K = r.J_meV / KB_MEV
results.append(r)
return results
# ---------------------------------------------------------------------------
# High-level convenience: read every config dir under <compound>/
# ---------------------------------------------------------------------------
[docs]
def read_all_energies(compound_dir: str, code: str,
max_index: Optional[int] = None,
scf_basename: Optional[str] = None
) -> Dict[int, EnergyResult]:
"""Walk ``<compound_dir>/configN/`` and run the energy parser on each.
Parameters
----------
compound_dir : str
Root directory containing the ``configN/`` subdirectories.
code : {'vasp', 'wien2k'}
Forwarded to :func:`mag4.energy.get_energy`.
max_index : int, optional
If given, only configs ``config1`` … ``configN`` are read.
Useful to short-circuit on very large compound dirs.
Returns
-------
dict[int, EnergyResult]
Map ``config_index → EnergyResult``. Empty if no ``configN/``
directory was found.
"""
out: Dict[int, EnergyResult] = {}
if not os.path.isdir(compound_dir):
return out
cfg_re = re.compile(r"^config(\d+)$")
for entry in sorted(os.listdir(compound_dir)):
m = cfg_re.match(entry)
if not m:
continue
idx = int(m.group(1))
if max_index is not None and idx > max_index:
continue
path = os.path.join(compound_dir, entry)
if os.path.isdir(path):
out[idx] = get_energy(path, code, scf_basename=scf_basename)
return out
# ---------------------------------------------------------------------------
# Output: write j_values.dat
# ---------------------------------------------------------------------------
[docs]
def write_j_values_dat(results: List[JResult], path: str,
compound_dir: str, code: str,
spin: Optional[float] = None,
supercell: Optional[Tuple[int, int, int]] = None,
energies: Optional[dict] = None,
config_sz: Optional[Dict[int, int]] = None,
mags: Optional[Dict[int, Optional[float]]] = None) -> None:
"""Persist the four-state results to a fixed-width text file.
Beyond the J table this records the **supercell** used and a
**per-configuration** block (total energy, SCF iterations, converged moment
vs the configuration's ``Sz``) so the file is self-contained — everything
shown on screen, plus the moment/Sz consistency check.
# mag4-extract — four-state J values
# compound: <dir> code: vasp spin: 3.5
# supercell: 1x1x2
#
# Per-configuration energies
# config iter energy (eV) Sz mag exp|mag| status
# ...
#
# name d_mean(A) JS2(eV) JS2(meV) J(meV) J(K) status
J1 3.8987 -0.12345678 -123.456789 -10.082 -116.99 OK
"""
width_name, width_d, width_e, width_meV = 6, 10, 14, 14
width_J, width_K = 10, 10
has_spin = spin is not None and spin > 0
mu = 2.0 * spin if has_spin else None # saturated |moment| = 2·S
def _fmt(x: Optional[float], w: int, fmt: str = "{:>{w}.6f}") -> str:
return fmt.format(x, w=w) if x is not None else "{:>{w}}".format("N/A", w=w)
with open(path, "w") as f:
f.write("# mag4-extract — four-state J values\n")
f.write(f"# compound: {compound_dir} code: {code}"
+ (f" spin: {spin}" if has_spin else "")
+ "\n")
if supercell is not None:
f.write(f"# supercell: {supercell[0]}x{supercell[1]}x{supercell[2]}\n")
if energies:
f.write("#\n# Per-configuration energies\n")
mag_hdr = (f" {'Sz':>4} {'mag':>9} {'exp|mag|':>9}"
if mags else "")
f.write(f"# {'config':>7} {'iter':>5} {'energy (eV)':>18}"
f"{mag_hdr} status\n")
for idx, er in sorted(energies.items()):
e_s = f"{er.energy_eV:.6f}" if er.energy_eV is not None else "—"
line = f"# config{idx:<3} {er.n_iter:>5} {e_s:>18}"
status = "OK" if er.converged else "FAILED"
if mags is not None:
sz = config_sz.get(idx) if config_sz else None
mg = mags.get(idx)
exp = abs(mu * sz) if (mu is not None and sz is not None) else None
line += (f" {(sz if sz is not None else '—')!s:>4} "
f"{(f'{mg:.3f}' if mg is not None else '—'):>9} "
f"{(f'{exp:.3f}' if exp is not None else '—'):>9}")
# flag a moment that disagrees with Sz (|mag| vs 2·S·|Sz|)
if mg is not None and exp is not None and abs(abs(mg) - exp) > 0.5:
status += " ⚠mag≠Sz"
f.write(line + f" {status}\n")
f.write("#\n")
header = (f"# {'name':<{width_name}} {'d_mean(A)':>{width_d}}"
f" {'JS2(eV)':>{width_e}} {'JS2(meV)':>{width_meV}}")
if has_spin:
header += f" {'J(meV)':>{width_J}} {'J(K)':>{width_K}}"
header += " status\n"
f.write(header)
for r in results:
line = (f" {r.name:<{width_name}} {r.distance_A:>{width_d}.4f}"
f" {_fmt(r.JS2_eV, width_e)} {_fmt(r.JS2_meV, width_meV)}")
if has_spin:
line += f" {_fmt(r.J_meV, width_J)} {_fmt(r.J_K, width_K)}"
line += " " + ("OK" if r.computed else (r.error or "missing"))
f.write(line + "\n")