"""Linear least-squares (energy-mapping) estimator for Heisenberg couplings.
Layer: **common** — the back-end-agnostic estimator behind
``mag4-extract --method energy-mapping``. The model energy of a magnetic
configuration is linear in the unknowns,
E(config) = E0 + Σ_k JS²_k · c_k(config),
where ``c_k`` are the integer coefficients written by
:func:`mag4.energy_mapping.write_energy_mapping_report` (Ising spins ``±1``, so
each fitted coupling is the ``JS²`` product; divide by ``S²`` for the bare
``J``). Because the model is linear, the RMS-optimal parameters come from a
single ``np.linalg.lstsq`` solve — the global solution, no iteration.
Ported from the standalone ``fit_J.py`` precursor. Before fitting, columns that
are indeterminate are dropped:
* **constant** columns (the same coefficient in every order) are collinear with
``E0`` and only shift it;
* columns flagged as **folding** onto a periodic self-image (unreliable);
* columns **exactly proportional** to another kept column (aliased in a
too-small cell) — only their combination is determined.
A near-collinearity (VIF) diagnostic and covariance-based error bars flag
couplings whose split is poorly conditioned even when the overall fit is tight.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple
import numpy as np
log = logging.getLogger(__name__)
[docs]
def spin_power(label: str) -> int:
"""Power of ``S`` by which a fitted coupling is divided to get the bare
exchange: pairwise enters as ``J·S²``, the four-spin ring term as
``Jring·S⁴`` (Ising spins ``±1``)."""
return 4 if "ring" in label.lower() else 2
[docs]
@dataclass
class FitResult:
"""Outcome of one energy-mapping fit (energies and JS² in **eV**)."""
e0: float
js2: Dict[str, float] # fitted JS² per kept coupling
se: Dict[str, float] # standard error of JS² (may be empty)
se_e0: Optional[float]
dropped: Dict[str, str] # label -> human-readable reason
kept: List[str]
vif: Dict[str, float]
rms_eV: float # RMS of residuals
rms_dof_eV: float # RMS per (n - n_param) DOF
n_data: int
n_param: int
residuals: Dict[int, float] # config number -> residual (eV)
max_corr: Optional[Tuple[str, str, float]]
def _vif(coef_matrix: np.ndarray, kept: List[str]) -> Dict[str, float]:
"""Variance-inflation factor per kept coupling (no intercept column)."""
n, k = coef_matrix.shape
out: Dict[str, float] = {}
for j in range(k):
y = coef_matrix[:, j]
others = [m for m in range(k) if m != j]
X = (np.column_stack([np.ones(n)] + [coef_matrix[:, m] for m in others])
if others else np.ones((n, 1)))
beta, *_ = np.linalg.lstsq(X, y, rcond=None)
ss_res = float(np.sum((y - X @ beta) ** 2))
ss_tot = float(np.sum((y - y.mean()) ** 2))
if ss_tot < 1e-30:
out[kept[j]] = 1.0
continue
r2 = max(0.0, min(1.0, 1.0 - ss_res / ss_tot))
out[kept[j]] = float("inf") if r2 > 1 - 1e-12 else 1.0 / (1.0 - r2)
return out
[docs]
def fit_couplings(records: List[dict], labels: List[str],
energies: Dict[int, float], *,
weight_by_deg: bool = False,
keep_collinear: bool = False,
unreliable: Sequence[str] = ()) -> FitResult:
"""Fit ``E = E0 + Σ JS²_k c_k`` by linear least-squares.
Parameters
----------
records : list of ``{"number", "coefficients": {label: c}, "degeneracy", "sz"}``.
labels : coupling names in column order.
energies : ``{config_number: energy_eV}`` for the configurations to fit
(caller filters to converged ones).
weight_by_deg : weight each residual by ``sqrt(degeneracy)``.
keep_collinear : keep every column (skip the indeterminate-column drops).
unreliable : coupling names known to fold onto self-images (dropped).
Raises
------
ValueError
If fewer usable energies than parameters (``E0 + kept couplings``).
"""
fit_recs = [r for r in records if r["number"] in energies]
cfg_order = [r["number"] for r in fit_recs]
n = len(cfg_order)
if n == 0:
raise ValueError("No configurations with a usable energy to fit.")
coef_of = {r["number"]: r["coefficients"] for r in fit_recs}
deg_of = {r["number"]: r.get("degeneracy", 1) for r in fit_recs}
# ---- decide which couplings to drop -------------------------------------
dropped: Dict[str, str] = {}
const_val: Dict[str, float] = {}
for lbl in labels:
col = [coef_of[c].get(lbl, 0) for c in cfg_order]
if max(col) - min(col) == 0:
const_val[lbl] = col[0]
geo_bad = [l for l in labels if l in set(unreliable)]
prop_drop: Dict[str, Tuple[str, float]] = {}
if keep_collinear:
kept = list(labels)
else:
surviving = [l for l in labels if l not in const_val and l not in geo_bad]
colcache = {l: np.array([coef_of[c].get(l, 0.0) for c in cfg_order], float)
for l in surviving}
kept = []
for lbl in surviving:
col = colcache[lbl]
match = None
for base in kept:
bcol = colcache[base]
nz = np.abs(bcol) > 1e-9
if not nz.any():
continue
alpha = col[nz][0] / bcol[nz][0]
if abs(alpha) > 1e-9 and np.allclose(col, alpha * bcol, atol=1e-9):
match = (base, alpha)
break
if match:
prop_drop[lbl] = match
else:
kept.append(lbl)
for lbl in labels:
if lbl in kept:
continue
if lbl in geo_bad and lbl not in const_val:
dropped[lbl] = "folds onto a periodic self-image (indeterminate)"
elif lbl in prop_drop:
base, alpha = prop_drop[lbl]
dropped[lbl] = (f"column = {alpha:+g}·{base} (aliased in this cell; "
f"folded into {base})")
elif lbl in const_val:
v = const_val[lbl]
dropped[lbl] = ("coefficient 0 in every order (never enters E)"
if abs(v) < 1e-12 else
f"constant coefficient {v:+g} (collinear with E0)")
n_param = 1 + len(kept)
if n < n_param:
raise ValueError(
f"Only {n} usable energ(ies) for {n_param} parameters "
f"(E0 + {len(kept)} coupling(s)); need ≥ {n_param}. Use a larger "
f"supercell to generate more inequivalent orders.")
# ---- weighted least squares ---------------------------------------------
A = np.array([[1.0] + [coef_of[c].get(l, 0.0) for l in kept]
for c in cfg_order], dtype=float)
b = np.array([energies[c] for c in cfg_order], dtype=float)
w = np.array([np.sqrt(deg_of[c]) if (weight_by_deg and deg_of[c]) else 1.0
for c in cfg_order], dtype=float)
Aw, bw = A * w[:, None], b * w
x, *_ = np.linalg.lstsq(Aw, bw, rcond=None)
e0 = float(x[0])
js2 = {l: float(x[1 + k]) for k, l in enumerate(kept)}
if np.linalg.matrix_rank(Aw) < A.shape[1]:
log.warning("Design matrix is rank-deficient after dropping constant "
"columns — some kept couplings remain collinear, so "
"individual JS² values may be unreliable.")
resid = A @ x - b
rms = float(np.sqrt(np.mean(resid ** 2)))
dof = max(n - n_param, 1)
rms_dof = float(np.sqrt(np.sum(resid ** 2) / dof))
residuals = {cfg_order[i]: float(resid[i]) for i in range(n)}
se: Dict[str, float] = {}
se_e0: Optional[float] = None
corr = None
try:
resid_w = Aw @ x - bw
sigma2 = float(resid_w @ resid_w) / dof
cov = sigma2 * np.linalg.inv(Aw.T @ Aw)
sd = np.sqrt(np.clip(np.diag(cov), 0.0, None))
se_e0 = float(sd[0])
for k, l in enumerate(kept):
se[l] = float(sd[1 + k])
d = np.sqrt(np.clip(np.diag(cov), 1e-300, None))
corr = cov / np.outer(d, d)
except np.linalg.LinAlgError:
pass
vif = _vif(A[:, 1:], kept) if len(kept) >= 2 else {}
max_corr: Optional[Tuple[str, str, float]] = None
if corr is not None and len(kept) >= 2:
cc = corr[1:, 1:]
worst, wi, wj = 0.0, None, None
for i in range(len(kept)):
for j in range(i + 1, len(kept)):
if abs(cc[i, j]) > worst:
worst, wi, wj = abs(cc[i, j]), i, j
if wi is not None:
max_corr = (kept[wi], kept[wj], float(worst))
return FitResult(
e0=e0, js2=js2, se=se, se_e0=se_e0, dropped=dropped, kept=kept,
vif=vif, rms_eV=rms, rms_dof_eV=rms_dof, n_data=n, n_param=n_param,
residuals=residuals, max_corr=max_corr,
)