"""Generate, group and merge the four spin configurations of the four-state method."""
from __future__ import annotations
import logging
import math
import random
from collections import defaultdict
from fractions import Fraction
from math import gcd
from typing import Dict, List, Optional, Tuple
import numpy as np
from scipy.spatial.distance import cdist
from mag4.constants import Dimer, FourStateConfig, next_config_idx
from mag4.lattice import get_cartesian_coords
log = logging.getLogger(__name__)
[docs]
def build_reference_bath(n_atoms: int,
user_ref: Optional[List[int]] = None) -> list:
"""Build the reference magnetic bath.
Default behaviour is fully ferromagnetic (all +1). Pass ``user_ref`` as a
list of ±1 of length ``n_atoms`` to override.
Raises
------
ValueError
If the length is wrong or values other than ±1 are present.
"""
if user_ref is not None:
if len(user_ref) != n_atoms:
raise ValueError(
f"--ref-state has {len(user_ref)} values "
f"but there are {n_atoms} magnetic atoms."
)
invalid = [v for v in user_ref if v not in (+1, -1)]
if invalid:
raise ValueError(
f"--ref-state values must be +1 or -1; got: {invalid}"
)
return list(user_ref)
return [+1 for _ in range(n_atoms)]
# ---------------------------------------------------------------------------
# Reference bath from a supercell MAGMOM-style string (--ref-magmom)
# ---------------------------------------------------------------------------
[docs]
def parse_magmom_tokens(s: str) -> List[float]:
"""Expand a VASP MAGMOM-style string into a flat list of floats.
Accepts plain values and ``n*v`` repetitions, whitespace- or
comma-separated: ``"32*0 1 -1 2*1.5"`` → ``[0.0]*32 + [1.0, -1.0, 1.5, 1.5]``.
"""
vals: List[float] = []
for tok in s.replace(",", " ").split():
if "*" in tok:
n_str, v_str = tok.split("*", 1)
n = int(n_str)
if n < 0:
raise ValueError(f"negative repeat count in '{tok}'")
vals.extend([float(v_str)] * n)
else:
vals.append(float(tok))
return vals
[docs]
def build_reference_from_magmom(magmom: str, crystal) -> list:
"""±1 reference bath from a MAGMOM-style string of the WORKING SUPERCELL.
The string is written exactly as one would set VASP's ``MAGMOM`` for the
supercell — POSCAR atom order (``crystal.all_species`` species blocks),
``n*v`` repetitions allowed — e.g. for a 2×2×1 La₂CuO₄ √2 cell::
32*0 1 1 1 -1 -1 -1 -1 1 -1 -1 -1 1 1 1 1 -1 64*0
Only the SIGNS of the magnetic entries are used (any magnitude, e.g. ±1 or
±0.6); non-magnetic entries must be 0 and are validated against the
species layout so an atom-order mismatch fails loudly. As a shorthand the
string may instead carry just the ``N_mag`` magnetic values (in the
magnetic-sublattice order). Returns the ±1 bath in
``crystal.target_species`` order.
"""
vals = parse_magmom_tokens(magmom)
all_atoms = [a for grp in crystal.all_species for a in grp]
n_all, n_mag = len(all_atoms), len(crystal.target_species)
if len(vals) == n_mag and n_mag != n_all:
picked = vals # magnetic-block shorthand
elif len(vals) == n_all:
mag_labels = {a[0] for a in crystal.target_species}
by_label = {}
for atom, v in zip(all_atoms, vals):
lbl = atom[0]
if lbl in mag_labels:
by_label[lbl] = v
elif abs(v) > 1e-8:
raise ValueError(
f"--ref-magmom: non-magnetic atom {lbl} carries a nonzero "
f"moment ({v:g}) — the string does not match the "
f"supercell's POSCAR atom order "
f"({' '.join(f'{crystal.elements[i]}×{len(g)}' for i, g in enumerate(crystal.all_species))}).")
picked = []
for a in crystal.target_species:
v = by_label.get(a[0], 0.0)
if abs(v) < 1e-8:
raise ValueError(
f"--ref-magmom: magnetic atom {a[0]} has moment 0 — every "
"magnetic site needs a signed value.")
picked.append(v)
else:
raise ValueError(
f"--ref-magmom has {len(vals)} value(s); expected {n_all} (full "
f"supercell, POSCAR order) or {n_mag} (magnetic atoms only).")
return [+1 if v > 0 else -1 for v in picked]
# ---------------------------------------------------------------------------
# Reference bath from requested coupling signs (--ref-couplings)
# ---------------------------------------------------------------------------
#
# The user states which couplings should be ferro- or antiferro-magnetic, e.g.
# "J1:AFM,J2:AFM". We then look for a collinear ±1 bath such that, for every
# bond of coupling Jk, the spin product s_i·s_j equals the requested sign
# (AFM → −1, FM → +1). Periodic images are included, so the bath is consistent
# across the (super)cell boundary.
#
# 1. Exact route — signed graph 2-colouring (BFS): seed +1, propagate
# s_j = s_i · sign_k. Succeeds iff the constraints are non-frustrated.
# 2. Fallback — Ising-energy minimisation when the constraints are frustrated
# (odd cycles, an atom bonded to its own image under AFM, an incommensurate
# cell, ...). Returns the best collinear compromise and reports which
# couplings could not be fully satisfied.
# AFM ⇒ antiparallel (s_i·s_j = −1); FM ⇒ parallel (+1).
_SIGN_WORDS = {
"AFM": -1, "AF": -1, "A": -1, "-": -1, "-1": -1,
"FM": +1, "F": +1, "+": +1, "+1": +1,
}
[docs]
def parse_ref_couplings(spec: str, coupling_names: List[str]) -> Dict[str, int]:
"""Parse ``"J1:AFM,J2:FM"`` → ``{"J1": -1, "J2": +1}``.
Raises ``ValueError`` on unknown coupling names or unrecognised signs.
"""
wanted: Dict[str, int] = {}
for item in spec.split(","):
item = item.strip()
if not item:
continue
if ":" not in item:
raise ValueError(
f"--ref-couplings entry '{item}' must look like 'J1:AFM'")
name, word = item.split(":", 1)
name, word = name.strip(), word.strip().upper()
if name not in coupling_names:
raise ValueError(
f"--ref-couplings: unknown coupling '{name}' "
f"(known: {', '.join(coupling_names)})")
if word not in _SIGN_WORDS:
raise ValueError(
f"--ref-couplings: '{word}' for {name} must be AFM or FM")
wanted[name] = _SIGN_WORDS[word]
if not wanted:
raise ValueError("--ref-couplings is empty")
return wanted
def _periodic_edges(coords: np.ndarray, lat_mat: np.ndarray,
distance: float, tol: float) -> List[Tuple[int, int]]:
"""All bonds (i, j) at ``distance`` ± ``tol``, periodic images included.
A bond between atom *i* and a periodic image of atom *j* still couples spins
``s_i`` and ``s_j`` (the image carries the same spin). Self-image bonds are
returned as ``(i, i)``. Each physical bond is counted once.
"""
n = len(coords)
cell_len = [float(np.linalg.norm(lat_mat[:, k])) for k in range(3)]
reach = max(1, int(math.ceil(distance / min(cell_len))))
seen: set = set()
edges: List[Tuple[int, int]] = []
for sa in range(-reach, reach + 1):
for sb in range(-reach, reach + 1):
for sc in range(-reach, reach + 1):
sh = (sa, sb, sc)
disp = lat_mat @ np.array(sh, dtype=float)
dmat = cdist(coords, coords + disp)
for i in range(n):
for j in range(n):
if abs(dmat[i, j] - distance) < tol:
mirror = (j, i, (-sa, -sb, -sc))
if (i, j, sh) in seen or mirror in seen:
continue
seen.add((i, j, sh))
edges.append((i, j))
return edges
def _signed_two_coloring(n: int, signed_edges: List[Tuple[int, int, int]]):
"""Exact signed 2-colouring.
``signed_edges`` are ``(i, j, want)`` with ``want = s_i·s_j`` desired.
Returns ``(colors, None)`` on success or ``(None, conflict)`` if frustrated,
where ``conflict`` is ``("self-image", i)`` or ``("cycle", (u, v))``.
"""
color = [0] * n
adj: Dict[int, list] = defaultdict(list)
for i, j, want in signed_edges:
if i == j:
if want == -1:
return None, ("self-image", i)
continue # FM self-image: trivially satisfied
adj[i].append((j, want))
adj[j].append((i, want))
for start in range(n):
if color[start] != 0:
continue
color[start] = +1
stack = [start]
while stack:
u = stack.pop()
for v, want in adj[u]:
need = color[u] * want
if color[v] == 0:
color[v] = need
stack.append(v)
elif color[v] != need:
return None, ("cycle", (u, v))
return [c if c != 0 else +1 for c in color], None
def _energy(s: list, weighted_edges: List[Tuple[int, int, int]]) -> int:
return sum(w * s[i] * s[j] for i, j, w in weighted_edges)
def _min_energy_bath(n: int, weighted_edges: List[Tuple[int, int, int]],
max_brute: int = 22, restarts: int = 16,
seed: int = 0) -> list:
"""Lowest-energy collinear bath for ``E = Σ w·s_i·s_j``.
``w = +1`` for AFM bonds (minimised antiparallel), ``−1`` for FM bonds.
``s_0`` is fixed to +1 (global spin-flip is degenerate); ties break toward
the most ferromagnetic bath. Exact brute force for ``n ≤ max_brute``,
otherwise greedy local search with random restarts.
"""
if n <= max_brute:
best_s, best_E = None, None
for bits in range(1 << (n - 1)):
s = [1] + [1 if (bits >> k) & 1 else -1 for k in range(n - 1)]
E = _energy(s, weighted_edges)
if best_E is None or E < best_E or (E == best_E and sum(s) > sum(best_s)):
best_E, best_s = E, s
return best_s
adj: Dict[int, list] = defaultdict(list)
for i, j, w in weighted_edges:
if i == j:
continue
adj[i].append((j, w))
adj[j].append((i, w))
rng = random.Random(seed)
best_s, best_E = None, None
for r in range(restarts):
s = [1] * n if r == 0 else [rng.choice((1, -1)) for _ in range(n)]
s[0] = 1
improved = True
while improved:
improved = False
for i in range(1, n):
dE = -2 * sum(w * s[i] * s[j] for j, w in adj[i])
if dE < 0:
s[i] = -s[i]
improved = True
E = _energy(s, weighted_edges)
if best_E is None or E < best_E or (E == best_E and sum(s) > sum(best_s)):
best_E, best_s = E, list(s)
return best_s
[docs]
def build_reference_from_couplings(atoms: list, lat_mat: np.ndarray,
couplings: list, spec: str,
tol: float) -> Tuple[list, dict]:
"""Build a reference bath that respects the requested coupling signs.
Parameters
----------
atoms : magnetic atoms in legacy ``[label, x, y, z, …]`` format.
lat_mat : 3×3 lattice matrix (columns are lattice vectors).
couplings : ``[[name, distance], …]``.
spec : e.g. ``"J1:AFM,J2:FM"``.
tol : distance tolerance for the bond search.
Returns ``(ref_bath, report)``. ``report`` carries the method used, any
frustration conflict, and per-coupling ``(satisfied, total)`` bond counts.
"""
names = [c[0] for c in couplings]
wanted = parse_ref_couplings(spec, names)
dist_of = {c[0]: c[1] for c in couplings}
coords = get_cartesian_coords(atoms, lat_mat)
n = len(atoms)
signed_edges: List[Tuple[int, int, int]] = []
weighted_edges: List[Tuple[int, int, int]] = []
per_coupling: Dict[str, List[Tuple[int, int]]] = {}
for name, want in wanted.items():
e = _periodic_edges(coords, lat_mat, dist_of[name], tol)
per_coupling[name] = e
w = +1 if want == -1 else -1 # AFM penalises parallel, FM penalises antiparallel
for i, j in e:
signed_edges.append((i, j, want))
weighted_edges.append((i, j, w))
color, conflict = _signed_two_coloring(n, signed_edges)
if color is not None:
ref, method = color, "exact (signed 2-colouring)"
else:
ref = _min_energy_bath(n, weighted_edges)
method = "approximate (frustrated → energy minimisation)"
coup_report: Dict[str, Tuple[int, int]] = {}
for name, want in wanted.items():
e = per_coupling[name]
sat = sum(1 for i, j in e if ref[i] * ref[j] == want)
coup_report[name] = (sat, len(e))
report = {
"method": method,
"conflict": conflict,
"wanted": wanted,
"couplings": coup_report,
}
return ref, report
# ---------------------------------------------------------------------------
# Reference bath from a propagation vector (--ref-kvector)
# ---------------------------------------------------------------------------
#
# For a commensurate collinear order with propagation vector k (in reciprocal
# lattice units of the *parent* cell), the spin of the atom at parent position
# R is s = sign(cos(2π k·R)). Each magnetic atom in the working (super)cell
# has supercell fractional coords f; its parent coords are R = f · (na,nb,nc).
#
# The cell hosts the order without a seam iff k_d·N_d is an integer on every
# axis (the supercell spans a whole number of magnetic periods).
[docs]
def parse_kvector(spec: str) -> List[Fraction]:
"""Parse ``"1/2 1/2 0"`` or ``"0.5,0.5,0"`` → ``[Fraction(1,2), …]``."""
toks = spec.replace(",", " ").split()
if len(toks) != 3:
raise ValueError(
f"--ref-kvector needs 3 components (e.g. '1/2 1/2 0'), got {len(toks)}")
try:
return [Fraction(t) for t in toks]
except (ValueError, ZeroDivisionError) as exc:
raise ValueError(f"--ref-kvector: cannot parse '{spec}' ({exc})")
def _commensurate_supercell(kf: List[Fraction],
supercell: Tuple[int, int, int]) -> Tuple[int, int, int]:
"""Smallest per-axis multiplier ≥ current that makes k_d·N_d integer."""
out = []
for d in range(3):
q = kf[d].denominator # magnetic period along d (parent cells)
cur = supercell[d]
out.append(cur * q // gcd(cur, q)) # lcm(cur, q)
return tuple(out)
[docs]
def build_reference_from_kvector(atoms: list, kfrac: List[Fraction],
supercell: Tuple[int, int, int],
node_tol: float = 1e-6) -> Tuple[list, dict]:
"""Build a collinear reference bath ``s_i = sign(cos(2π k·R_i))``.
Parameters
----------
atoms : magnetic atoms in legacy ``[label, x, y, z, …]`` format, with
fractional coordinates in the working (super)cell.
kfrac : propagation vector (3 ``Fraction``) in parent reciprocal-lattice units.
supercell : ``(na, nb, nc)`` multipliers used to build the working cell.
Returns ``(ref_bath, report)``.
"""
na, nb, nc = supercell
kf = [Fraction(k) for k in kfrac]
kflo = [float(k) for k in kf]
# Reference the phase to the first magnetic atom (a global spin-flip is
# physically irrelevant). This puts atom 0 at an antinode and makes the
# commensurate phase differences clean integers/half-integers — avoiding
# spurious "node" hits when every site happens to sit at k·R = ¼ in absolute
# terms (e.g. a layer at z = ½ with k = (0,0,½)).
if atoms:
r0 = (atoms[0][1] * na, atoms[0][2] * nb, atoms[0][3] * nc)
else:
r0 = (0.0, 0.0, 0.0)
ref: list = []
node_atoms: List[str] = []
for a in atoms:
rp = (a[1] * na - r0[0], a[2] * nb - r0[1], a[3] * nc - r0[2])
phase = 2.0 * math.pi * (kflo[0]*rp[0] + kflo[1]*rp[1] + kflo[2]*rp[2])
cval = math.cos(phase)
if abs(cval) < node_tol: # genuine relative node → ambiguous
node_atoms.append(a[0])
ref.append(1 if cval >= 0 else -1)
kN = [kf[d] * supercell[d] for d in range(3)]
incommensurate = [("a", "b", "c")[d] for d in range(3) if kN[d].denominator != 1]
report = {
"method": f"propagation vector k=({_fmt_kvec(kf)})",
"kvec": tuple(kf),
"supercell": (na, nb, nc),
"incommensurate": incommensurate,
"suggest_supercell": _commensurate_supercell(kf, (na, nb, nc)),
"node_atoms": node_atoms,
}
return ref, report
def _fmt_kvec(kf: List[Fraction]) -> str:
def one(x: Fraction) -> str:
return str(x.numerator) if x.denominator == 1 else f"{x.numerator}/{x.denominator}"
return " ".join(one(x) for x in kf)
[docs]
def generate_four_states(ref_bath: list, dimer: Dimer,
coupling_name: str) -> List[FourStateConfig]:
"""Generate the four spin configurations ``uu, dd, ud, du`` for one dimer.
Only the two atoms of the dimer differ between configurations; the rest of
the magnetic bath is frozen to ``ref_bath``.
"""
configs = []
for label, s_a, s_b in [("uu", +1, +1), ("dd", -1, -1),
("ud", +1, -1), ("du", -1, +1)]:
spins = list(ref_bath)
spins[dimer.idx_a] = s_a
spins[dimer.idx_b] = s_b
configs.append(FourStateConfig(
idx = next_config_idx(),
coupling_name = coupling_name,
state_label = label,
spin_vector = tuple(spins),
dimer = dimer,
))
return configs
[docs]
def find_unique_configs(all_configs: List[FourStateConfig],
sym_ops: List[Tuple[int, ...]],
*,
time_reversal: bool = True
) -> List[List[FourStateConfig]]:
"""Group configurations equivalent under the magnetic (grey) group.
The group is the unitary crystal operations ``sym_ops`` plus, when
``time_reversal`` is True (the default), each of them composed with time
reversal T (global spin flip) — exact in a collinear calculation without
spin-orbit coupling. Returns a list of groups; the energy of any member of
a group can be used interchangeably in the J extraction formulas.
``time_reversal=False`` restricts to the unitary subgroup — used by
``--check-degeneracy`` to keep the T-redundant configurations as separate
DFT runs whose energy agreement then *tests* the antiunitary symmetry
(and the completeness of the spin model).
"""
def equiv(sv1, sv2):
sv2_rev = tuple(-s for s in sv2)
for op in sym_ops:
mapped = tuple(sv1[op[i]] for i in range(len(sv1)))
if mapped == sv2 or (time_reversal and mapped == sv2_rev):
return True
return False
groups: List[List[FourStateConfig]] = []
for cfg in all_configs:
matched = False
for group in groups:
if equiv(group[0].spin_vector, cfg.spin_vector):
group.append(cfg)
matched = True
break
if not matched:
groups.append([cfg])
return groups
[docs]
def time_reversal_pairs(unitary_groups: List[List[FourStateConfig]],
sym_ops: List[Tuple[int, ...]]
) -> List[Tuple[int, int]]:
"""Pairs of unitary-only classes that time reversal proves degenerate.
``unitary_groups`` must come from :func:`find_unique_configs` with
``time_reversal=False``. Two classes pair up when an ANTIUNITARY
operation P×T (unitary op composed with global spin flip) maps one
representative onto the other — their DFT energies must then coincide
even though no unitary operation relates them (they typically differ in
total Sz). Returns 1-based ``(uid_a, uid_b)`` index pairs; since the
unitary subgroup has index 2 in the grey group, each class belongs to at
most one pair.
"""
def tr_equiv(sv1, sv2):
sv2_rev = tuple(-s for s in sv2)
return any(tuple(sv1[op[i]] for i in range(len(sv1))) == sv2_rev
for op in sym_ops)
pairs: List[Tuple[int, int]] = []
used: set = set()
for a in range(len(unitary_groups)):
if a in used:
continue
for b in range(a + 1, len(unitary_groups)):
if b in used:
continue
if tr_equiv(unitary_groups[a][0].spin_vector,
unitary_groups[b][0].spin_vector):
pairs.append((a + 1, b + 1))
used.add(a)
used.add(b)
break
return pairs