"""Energy-mapping (linear least-squares) configuration generator.
Layer: **common** — an alternative to the four-state method for extracting
Heisenberg *J* from DFT total energies. Where the four-state method targets a
handful of ``uu/dd/ud/du`` configurations per bond, energy-mapping enumerates
**every symmetry-/coefficient-inequivalent collinear order** of a given
(super)cell and later fits
E(config) = E0 + Σ_k J_k · c_k(config)
by a single linear least-squares solve over all configurations (see
``mag4-extract --method energy-mapping``, Phase 3).
This module produces, for one (super)cell, the list of inequivalent orders and
their integer coupling coefficients ``c_k`` (the design-matrix rows). The
``±1`` Ising enumeration and the coefficient counting mirror the standalone
``ising_magnetic.py`` precursor, re-expressed on mag4's :class:`CrystalData`.
"""
from __future__ import annotations
import itertools
import logging
from typing import Dict, List, Optional, Tuple
import numpy as np
from scipy.spatial.distance import cdist
from mag4.constants import CrystalData, FourStateConfig
from mag4.lattice import frac_to_cart
log = logging.getLogger(__name__)
#: Above this magnetic-atom count the 2^N enumeration is rejected outright
#: (overridable per call). 2^20 ≈ 1e6 orders is already minutes of work.
DEFAULT_MAX_ATOMS = 20
#: Log a slow-enumeration warning at and above this many magnetic atoms.
_WARN_ATOMS = 16
def _cartesian_with_spin(atoms: list, lat_mat: np.ndarray) -> list:
"""``[label, x, y, z, spin]`` in Cartesian coordinates (spin slot = 0)."""
frac = np.array([[a[1], a[2], a[3]] for a in atoms], dtype=float)
cart = frac_to_cart(frac, lat_mat)
return [[atoms[i][0], cart[i, 0], cart[i, 1], cart[i, 2], 0]
for i in range(len(atoms))]
def _image_shell(cell_atoms: list, lat_mat: np.ndarray) -> list:
"""Replicate ``cell_atoms`` over the 3×3×3 block of periodic images.
Each image keeps its **parent** label, so the spin map resolves an image
to the spin of its in-cell parent.
"""
shifts = [(x, y, z) for x in (-1, 0, 1) for y in (-1, 0, 1) for z in (-1, 0, 1)]
names = [a[0] for a in cell_atoms]
coords = np.array([[a[1], a[2], a[3]] for a in cell_atoms], dtype=float)
out: list = []
for s in shifts:
disp = lat_mat @ np.array(s, dtype=float)
shifted = coords + disp
for i in range(len(names)):
out.append([names[i], shifted[i, 0], shifted[i, 1], shifted[i, 2]])
return out
def _identify_bonds(cell_atoms: list, image_atoms: list, couplings: list,
tol: float, distmax: float) -> List[Tuple[str, str, str]]:
"""All (cell-atom, image-atom, coupling) bonds within ``d ± tol`` (< distmax)."""
cc = np.array([[a[1], a[2], a[3]] for a in cell_atoms], dtype=float)
ic = np.array([[a[1], a[2], a[3]] for a in image_atoms], dtype=float)
dmat = cdist(cc, ic)
bonds: List[Tuple[str, str, str]] = []
for name, dist in couplings:
mask = (dmat > dist - tol) & (dmat < dist + tol) & (dmat < distmax)
for i, j in np.argwhere(mask):
bonds.append((cell_atoms[i][0], image_atoms[j][0], name))
return bonds
def _coefficients(spin_map: Dict[str, int], bonds: list, couplings: list) -> tuple:
"""Integer coupling coefficients ``(c_1, …, c_n, |Sz|)`` for one ordering.
``c_k = ½ Σ_{⟨ij⟩∈k} s_i s_j`` (the ½ removes the double counting of the
cell↔image bond list, mirroring the four-state convention).
"""
coefs = []
for name, _ in couplings:
total = sum(spin_map.get(i, 0) * spin_map.get(j, 0)
for i, j, nm in bonds if nm == name)
coefs.append(total // 2)
coefs.append(abs(sum(spin_map.values())))
return tuple(coefs)
[docs]
def config_energy_coefficients(atoms: list, lat_mat: np.ndarray, couplings: list,
spin_vectors: List[Tuple[int, ...]], *,
tol: float = 5e-4,
distmax: Optional[float] = None
) -> List[Tuple[Dict[str, int], int]]:
"""Coupling coefficients of each spin order's energy.
For every spin vector (``±1`` per magnetic atom, in ``atoms`` order) returns
``({Jname: c_k}, |Sz|)`` with ``c_k = Σ_{⟨ij⟩∈k} s_i s_j`` summed over the
**unique** bonds of shell ``k`` in the cell — i.e. the design-matrix row, so
that ``E[config] = E0 + Σ_k c_k · (J_k · S²)``. This is what lets one verify
that a four-state combination isolates a single ``J`` (the other couplings'
coefficients cancel). Bonds are counted once via the same cell↔image folding
as :func:`_coefficients`.
"""
lat = np.asarray(lat_mat, dtype=float)
if distmax is None:
distmax = max(d for _, d in couplings) + max(tol, 1e-3)
cell_atoms = _cartesian_with_spin(atoms, lat)
image_atoms = _image_shell(cell_atoms, lat)
bonds = _identify_bonds(cell_atoms, image_atoms, couplings, tol, distmax)
names = [a[0] for a in cell_atoms]
out: List[Tuple[Dict[str, int], int]] = []
for sv in spin_vectors:
spin_map = {names[i]: int(sv[i]) for i in range(len(names))}
c = _coefficients(spin_map, bonds, couplings) # (c_1,…,c_n,|Sz|)
out.append(({couplings[k][0]: int(c[k]) for k in range(len(couplings))},
int(c[-1])))
return out
# ---------------------------------------------------------------------------
# Four-spin cyclic (ring) exchange (--jring)
# ---------------------------------------------------------------------------
def _identify_plaquettes(cell_atoms: list, image_atoms: list, nn_dist: float,
tol_pos: float = 0.15):
"""Nearest-neighbour square plaquettes ``(i, j, k, l)`` for the ring term.
A plaquette is a square: a cell atom *i* plus two NN neighbours *j, l* at
~90°, with the fourth corner ``k = i + (j−i) + (l−i)``. Each physical
square is found once from every corner (4× redundant), so
:func:`count_ring` divides by 4. Corners carry their *parent* atom name so
the spin map resolves images. Returns ``(plaquettes, n_unique, n_folded)``.
"""
rc = np.array([[a[1], a[2], a[3]] for a in cell_atoms], dtype=float)
rs = np.array([[a[1], a[2], a[3]] for a in image_atoms], dtype=float)
ns = [a[0] for a in image_atoms]
nc = [a[0] for a in cell_atoms]
plaquettes: list = []
folded = 0
for ci in range(len(cell_atoms)):
ri = rc[ci]
d = np.linalg.norm(rs - ri, axis=1)
nbr = np.where((d > nn_dist - tol_pos) & (d < nn_dist + tol_pos))[0]
for x in range(len(nbr)):
for y in range(x + 1, len(nbr)):
ja, jb = int(nbr[x]), int(nbr[y])
va, vb = rs[ja] - ri, rs[jb] - ri
denom = np.linalg.norm(va) * np.linalg.norm(vb)
if denom == 0:
continue
if abs(float(np.dot(va, vb) / denom)) > 0.5: # keep ~90°, drop 180°
continue
rk = ri + va + vb
hit = np.where(np.linalg.norm(rs - rk, axis=1) < tol_pos)[0]
if hit.size == 0:
continue
names = (nc[ci], ns[ja], ns[int(hit[0])], ns[jb])
plaquettes.append(names)
if len(set(names)) < 4:
folded += 1
return plaquettes, len(plaquettes) // 4, folded // 4
[docs]
def count_ring(plaquettes: list, spin_map: Dict[str, int]) -> int:
"""Ring coefficient = ``Σ s_i s_j s_k s_l`` over the (4×-redundant) squares,
divided by 4."""
total = sum(spin_map[i] * spin_map[j] * spin_map[k] * spin_map[l]
for (i, j, k, l) in plaquettes)
return total // 4
[docs]
def select_ring_plaquettes(couplings: list, cell_atoms: list, image_atoms: list,
prefer_label: Optional[str] = None,
prefer_dist: Optional[float] = None,
tol_pos: float = 0.15):
"""Pick the coupling shell that forms the square lattice for the ring term.
The shortest M–M distance is not necessarily the square edge (e.g. SrFeO₂'s
shortest Fe–Fe is a linear c-axis contact that forms no square). If the
user pins a shell (``prefer_label`` / ``prefer_dist``) it is used; otherwise
every shell is tried from shortest to longest and the first that closes into
squares is returned. Returns ``(plaquettes, n_unique, n_folded, dist)`` or
``(None, 0, 0, None)``.
"""
if prefer_dist is not None:
trials = [("J?", float(prefer_dist))]
elif prefer_label is not None:
match = [(n, d) for n, d in couplings if n == prefer_label]
if not match:
log.warning("--jring-coupling %s not among couplings %s; "
"falling back to the automatic shell scan.",
prefer_label, [n for n, _ in couplings])
trials = sorted(((n, d) for n, d in couplings), key=lambda x: x[1])
else:
trials = match
else:
trials = sorted(((n, d) for n, d in couplings), key=lambda x: x[1])
for name, dist in trials:
plaq, n_uni, folded = _identify_plaquettes(cell_atoms, image_atoms, dist, tol_pos)
if plaq:
return plaq, n_uni, folded, dist
log.info("--jring: shell %s (%.4f Å) forms no square plaquette "
"(likely a linear contact); trying the next shell.", name, dist)
return None, 0, 0, None
[docs]
def enumerate_orders(crystal: CrystalData, couplings: list,
tol: float = 5e-4, distmax: Optional[float] = None,
*, max_atoms: int = DEFAULT_MAX_ATOMS,
ring: bool = False, ring_coupling: Optional[str] = None,
ring_dist: Optional[float] = None, ring_tol: float = 0.15
) -> Tuple[List[dict], List[str], Optional[float]]:
"""Enumerate the inequivalent collinear magnetic orders of ``crystal``.
Every ``2^N`` ``±1`` assignment over the ``N`` magnetic atoms is grouped by
its coupling-coefficient signature ``(c_1, …, c_n, |Sz|)``; configurations
sharing a signature are energy-degenerate in the model, so one
representative ordering is kept per signature.
With ``ring=True`` a four-spin cyclic (ring) exchange coefficient is also
computed per order and a ``"Jring"`` column appended (see
:func:`select_ring_plaquettes`); the square-lattice shell is auto-detected
unless pinned with ``ring_coupling`` / ``ring_dist``.
Returns ``(records, labels, ring_distance)`` where ``labels`` are the
coupling names in column order (``"Jring"`` last when present),
``ring_distance`` is the square edge length (or ``None``), and each record
is::
{"number": int, # 1-based, deterministic
"coefficients": {name: c},
"sz": int, # |Sz| of the order
"degeneracy": int, # how many ±1 orders share this signature
"spin_vector": [±1, …]} # representative, in target_species order
Raises
------
ValueError
When ``N`` exceeds ``max_atoms`` (the ``2^N`` enumeration would be
intractable) — use a smaller supercell or raise ``max_atoms``.
"""
atoms = crystal.target_species
n = len(atoms)
if n > max_atoms:
raise ValueError(
f"Energy-mapping enumerates 2^{n} = {2**n:,} spin orders for {n} "
f"magnetic atoms, above the 2^{max_atoms} guard. Use a smaller "
f"supercell, or pass a larger max_atoms if you really mean it.")
if n >= _WARN_ATOMS:
log.warning("Enumerating 2^%d = %s spin orders — this may take a while.",
n, f"{2**n:,}")
if distmax is None:
distmax = max(d for _, d in couplings) + max(tol, 1e-3)
lat = crystal.lattice_matrix
cell_atoms = _cartesian_with_spin(atoms, lat)
image_atoms = _image_shell(cell_atoms, lat)
bonds = _identify_bonds(cell_atoms, image_atoms, couplings, tol, distmax)
labels = [c[0] for c in couplings]
# Optional four-spin ring term.
plaquettes = None
ring_distance: Optional[float] = None
if ring:
plaquettes, n_uni, n_folded, ring_distance = select_ring_plaquettes(
couplings, cell_atoms, image_atoms,
prefer_label=ring_coupling, prefer_dist=ring_dist, tol_pos=ring_tol)
if plaquettes is None:
log.warning("--jring: no coupling shell forms square plaquettes "
"(all are linear contacts); the ring term is skipped.")
else:
log.info("--jring: ring term on the %.4f Å square lattice "
"(%d square(s) per cell).", ring_distance, n_uni)
if n_folded:
log.warning("--jring: %d square(s) fold onto a repeated site in "
"this supercell — the ring coefficient may be "
"unreliable; enlarge the supercell.", n_folded)
has_ring = plaquettes is not None
if has_ring:
labels = labels + ["Jring"]
degeneracy: Dict[tuple, int] = {}
representative: Dict[tuple, list] = {}
names = [a[0] for a in cell_atoms]
for spin in itertools.product((1, -1), repeat=n):
spin_map = {names[i]: spin[i] for i in range(n)}
sig = _coefficients(spin_map, bonds, couplings) # (c_1, …, c_n, |Sz|)
if has_ring:
rc = count_ring(plaquettes, spin_map)
sig = sig[:-1] + (rc, sig[-1]) # (c_1, …, c_n, ring, |Sz|)
degeneracy[sig] = degeneracy.get(sig, 0) + 1
if sig not in representative:
representative[sig] = list(spin)
records: List[dict] = []
for number, sig in enumerate(sorted(degeneracy), 1):
records.append({
"number": number,
"coefficients": {labels[k]: int(sig[k]) for k in range(len(labels))},
"sz": int(sig[-1]),
"degeneracy": degeneracy[sig],
"spin_vector": representative[sig],
})
log.info("Energy-mapping: %d inequivalent order(s) out of 2^%d = %s.",
len(records), n, f"{2**n:,}")
return records, labels, ring_distance
[docs]
def build_unique_groups(records: List[dict]) -> List[List[FourStateConfig]]:
"""Wrap each record as a one-member group for the VASP/WIEN2k writers.
:func:`mag4.vasp.build_magmom_lines` / :func:`mag4.vasp.create_config_dirs`
consume ``List[List[FourStateConfig]]`` and read only ``spin_vector`` (for
the MAGMOM) plus ``coupling_name`` / ``state_label`` (for the printed
label). Energy-mapping has no per-bond label, so the coefficient signature
is used for display.
"""
groups: List[List[FourStateConfig]] = []
for rec in records:
coef_str = " ".join(f"{c:+d}*{name}"
for name, c in rec["coefficients"].items())
groups.append([FourStateConfig(
idx=rec["number"],
coupling_name="config",
state_label=f"deg{rec['degeneracy']},Sz{rec['sz']}",
spin_vector=tuple(rec["spin_vector"]),
dimer=None,
)])
rec["_coef_str"] = coef_str
return groups
[docs]
def write_energy_mapping_report(path: str, records: List[dict], labels: List[str],
couplings: list, supercell: Tuple[int, int, int],
name: str, *,
tr_pairs: Optional[List[tuple]] = None,
config_msgs: Optional[List] = None) -> None:
"""Write the design-matrix report consumed by ``mag4-extract``.
The format is deliberately simple and machine-parseable: a ``Method:`` tag,
the coupling labels + distances, then one whitespace-separated row per
configuration giving ``config<N>``, degeneracy, Sz and the integer
coefficient of each coupling (column order = ``labels``).
``tr_pairs`` (from ``--check-degeneracy``) adds one ``pair: configA ==
configB`` line per time-reversal-degenerate pair; ``mag4-extract`` reads
them back and reports ``|ΔE|`` per pair. ``config_msgs`` (magnetic Shubnikov
groups) are recorded as comment lines — informational, ignored by parsers.
"""
dist_of = {c[0]: c[1] for c in couplings}
na, nb, nc = supercell
width = max(6, max((len(l) for l in labels), default=6))
with open(path, "w") as f:
f.write(f"# mag4 energy-mapping report — {name}\n")
f.write("Method: energy-mapping\n")
f.write(f"Supercell: {na} {nb} {nc}\n")
f.write("Couplings: "
+ ", ".join(f"{l}={dist_of.get(l, float('nan')):.4f}" for l in labels)
+ "\n")
f.write(f"N_configs: {len(records)}\n")
f.write("# Model: E(config) = E0 + " +
" + ".join(f"{l}*c_{l}" for l in labels) + "\n")
header = (f"{'config':>10} {'deg':>6} {'Sz':>5} "
+ " ".join(f"{l:>{width}}" for l in labels))
f.write("#" + header[1:] + "\n")
for rec in records:
row = (f"{'config' + str(rec['number']):>10} "
f"{rec['degeneracy']:>6} {rec['sz']:>5} "
+ " ".join(f"{rec['coefficients'][l]:>{width}d}" for l in labels))
f.write(row + "\n")
if config_msgs is not None and any(m is not None for m in config_msgs):
f.write("# Magnetic (Shubnikov) group per order, BNS number-type "
"and point group (primed ops included):\n")
for rec, m in zip(records, config_msgs):
if m is not None:
pg = f", MPG {m.point_group()}" if m.mpg_symbol else ""
f.write(f"# config{rec['number']} MSG: {m.bns_number}-"
f"{m.type_roman} (parent {m.parent_symbol}{pg})\n")
if tr_pairs:
f.write("# Time-reversal degeneracy check (--check-degeneracy):\n")
f.write("# each pair is related only by the ANTIUNITARY symmetry\n")
f.write("# (global spin flip); the two DFT energies must coincide.\n")
for a, b in tr_pairs:
f.write(f"pair: config{a} == config{b}\n")
f.write(f"# Sum of degeneracies: {sum(r['degeneracy'] for r in records)}\n")
[docs]
def detect_report_method(report_path: str) -> str:
"""Return the ``Method:`` tag of a mag4 report (``four-state`` if absent).
Legacy four-state reports carry no ``Method:`` line, so their absence is
treated as ``four-state``.
"""
try:
with open(report_path) as f:
for _ in range(20):
line = f.readline()
if not line:
break
s = line.strip()
if s.lower().startswith("method:"):
return s.split(":", 1)[1].strip()
except FileNotFoundError:
pass
return "four-state"
[docs]
def parse_energy_mapping_report(report_path: str) -> dict:
"""Parse the design-matrix report written by :func:`write_energy_mapping_report`.
Returns ``{"method", "labels", "dist_of", "supercell", "records"}`` where
each record is ``{"number", "degeneracy", "sz", "coefficients": {label: c}}``.
"""
method: Optional[str] = None
labels: List[str] = []
dist_of: Dict[str, float] = {}
supercell: Tuple[int, int, int] = (1, 1, 1)
records: List[dict] = []
with open(report_path) as f:
for raw in f:
s = raw.strip()
if s.lower().startswith("method:"):
method = s.split(":", 1)[1].strip()
elif s.lower().startswith("supercell:"):
parts = s.split(":", 1)[1].split()
if len(parts) == 3:
supercell = (int(parts[0]), int(parts[1]), int(parts[2]))
elif s.lower().startswith("couplings:"):
for tok in s.split(":", 1)[1].split(","):
tok = tok.strip()
if "=" in tok:
name, dist = tok.split("=", 1)
name = name.strip()
labels.append(name)
try:
dist_of[name] = float(dist)
except ValueError:
dist_of[name] = float("nan")
elif s.startswith("config") and not s.startswith("#"):
cols = s.split()
if len(cols) < 3 + len(labels):
continue
num = int(cols[0][len("config"):])
coeffs = {labels[k]: int(cols[3 + k]) for k in range(len(labels))}
records.append({
"number": num,
"degeneracy": int(cols[1]),
"sz": int(cols[2]),
"coefficients": coeffs,
})
if method is None:
raise ValueError(
f"{report_path} has no 'Method:' line — not a mag4 energy-mapping "
f"report.")
return {
"method": method,
"labels": labels,
"dist_of": dist_of,
"supercell": supercell,
"records": records,
}