use crate::autodiff::optimizer::{Adam, LBfgs, Optimizer, OptimizerKind, Sgd};
use crate::autodiff::tape::{Tape, Var};
use crate::constants::MU_0;
use crate::error::{dimension_mismatch, invalid_param, Result};
use crate::vector3::Vector3;
#[derive(Debug, Clone)]
pub struct SpinConfig {
pub theta: Vec<f64>,
pub phi: Vec<f64>,
}
impl SpinConfig {
pub fn new(n: usize) -> Self {
Self {
theta: vec![0.0; n],
phi: vec![0.0; n],
}
}
pub fn random(n: usize, rng_seed: u64) -> Self {
let mut state: u64 = if rng_seed == 0 {
0xDEAD_BEEF_CAFE_BABE
} else {
rng_seed
};
let mut draw_unit = || {
state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1_442_695_040_888_963_407);
let bits = state >> 11;
(bits as f64) / ((1u64 << 53) as f64)
};
let mut theta = Vec::with_capacity(n);
let mut phi = Vec::with_capacity(n);
for _ in 0..n {
let u = draw_unit();
let v = draw_unit();
theta.push((1.0 - 2.0 * u).acos());
phi.push(2.0 * std::f64::consts::PI * v);
}
Self { theta, phi }
}
pub fn to_vectors(&self) -> Vec<Vector3<f64>> {
self.theta
.iter()
.zip(self.phi.iter())
.map(|(&t, &p)| Vector3::new(t.sin() * p.cos(), t.sin() * p.sin(), t.cos()))
.collect()
}
pub fn from_vectors(v: &[Vector3<f64>]) -> Result<Self> {
let mut theta = Vec::with_capacity(v.len());
let mut phi = Vec::with_capacity(v.len());
for vi in v {
let mag = vi.magnitude();
if mag < 1e-12 {
return Err(invalid_param(
"spin vector",
"zero-magnitude vector cannot be converted to (θ, φ)",
));
}
let nz = vi.z / mag;
let nz_clamped = nz.clamp(-1.0, 1.0);
theta.push(nz_clamped.acos());
phi.push(vi.y.atan2(vi.x));
}
Ok(Self { theta, phi })
}
pub fn params_flat(&self) -> Vec<f64> {
let n = self.theta.len();
let mut v = Vec::with_capacity(2 * n);
for i in 0..n {
v.push(self.theta[i]);
v.push(self.phi[i]);
}
v
}
pub fn set_params(&mut self, flat: &[f64]) -> Result<()> {
let n = self.theta.len();
if flat.len() != 2 * n {
return Err(dimension_mismatch(
&format!("{} params", 2 * n),
&format!("{} params", flat.len()),
));
}
for i in 0..n {
self.theta[i] = flat[2 * i];
self.phi[i] = flat[2 * i + 1];
}
Ok(())
}
pub fn n_spins(&self) -> usize {
self.theta.len()
}
}
#[derive(Debug, Clone)]
pub struct EnergyFunctional {
pub exchange_j: f64,
pub anisotropy_k: f64,
pub h_ext: [f64; 3],
pub ms: f64,
pub bonds: Vec<(usize, usize)>,
}
impl EnergyFunctional {
pub fn new(exchange_j: f64, anisotropy_k: f64, h_ext: [f64; 3], ms: f64) -> Self {
Self {
exchange_j,
anisotropy_k,
h_ext,
ms,
bonds: Vec::new(),
}
}
pub fn with_bonds(mut self, bonds: Vec<(usize, usize)>) -> Self {
self.bonds = bonds;
self
}
pub fn with_chain_bonds(mut self, n_spins: usize) -> Self {
let mut bonds = Vec::with_capacity(n_spins.saturating_sub(1));
for i in 0..n_spins.saturating_sub(1) {
bonds.push((i, i + 1));
}
self.bonds = bonds;
self
}
pub fn with_ring_bonds(mut self, n_spins: usize) -> Self {
let mut bonds = Vec::with_capacity(n_spins);
for i in 0..n_spins.saturating_sub(1) {
bonds.push((i, i + 1));
}
if n_spins >= 2 {
bonds.push((n_spins - 1, 0));
}
self.bonds = bonds;
self
}
pub fn energy<'t>(&self, tape: &'t Tape, theta_phi: &[Var<'t>]) -> Result<Var<'t>> {
if theta_phi.len() % 2 != 0 {
return Err(invalid_param(
"theta_phi",
"length must be even ([θ_0, φ_0, θ_1, φ_1, …])",
));
}
let n = theta_phi.len() / 2;
let mut mx = Vec::with_capacity(n);
let mut my = Vec::with_capacity(n);
let mut mz = Vec::with_capacity(n);
for i in 0..n {
let t = theta_phi[2 * i];
let p = theta_phi[2 * i + 1];
let sin_t = t.sin();
mx.push(sin_t * p.cos());
my.push(sin_t * p.sin());
mz.push(t.cos());
}
let mut energy = Var::leaf(tape, 0.0);
if self.exchange_j != 0.0 && !self.bonds.is_empty() {
let mut ex_sum = Var::leaf(tape, 0.0);
for &(i, j) in &self.bonds {
if i >= n || j >= n {
return Err(dimension_mismatch(
&format!("indices < {n}"),
&format!("bond ({i}, {j})"),
));
}
let dot = mx[i] * mx[j] + my[i] * my[j] + mz[i] * mz[j];
ex_sum = ex_sum + dot;
}
energy = energy + ex_sum * (-self.exchange_j);
}
if self.anisotropy_k != 0.0 {
let mut aniso_sum = Var::leaf(tape, 0.0);
for mz_i in mz.iter().take(n) {
aniso_sum = aniso_sum + *mz_i * *mz_i;
}
energy = energy + aniso_sum * (-self.anisotropy_k);
}
let zeeman_coeff = -MU_0 * self.ms;
if zeeman_coeff != 0.0
&& (self.h_ext[0] != 0.0 || self.h_ext[1] != 0.0 || self.h_ext[2] != 0.0)
{
let mut zee_sum = Var::leaf(tape, 0.0);
for i in 0..n {
let dot = mx[i] * self.h_ext[0] + my[i] * self.h_ext[1] + mz[i] * self.h_ext[2];
zee_sum = zee_sum + dot;
}
energy = energy + zee_sum * zeeman_coeff;
}
Ok(energy)
}
pub fn energy_f64(&self, config: &SpinConfig) -> f64 {
let n = config.n_spins();
let vs: Vec<Vector3<f64>> = config.to_vectors();
let mut e = 0.0_f64;
if self.exchange_j != 0.0 {
let mut s = 0.0_f64;
for &(i, j) in &self.bonds {
if i < n && j < n {
s += vs[i].dot(&vs[j]);
}
}
e -= self.exchange_j * s;
}
if self.anisotropy_k != 0.0 {
let mut s = 0.0_f64;
for vi in &vs {
s += vi.z * vi.z;
}
e -= self.anisotropy_k * s;
}
let zee_coeff = -MU_0 * self.ms;
if zee_coeff != 0.0 {
let mut s = 0.0_f64;
for vi in &vs {
s += vi.x * self.h_ext[0] + vi.y * self.h_ext[1] + vi.z * self.h_ext[2];
}
e += zee_coeff * s;
}
e
}
}
#[derive(Debug, Clone)]
pub struct StructureOptResult {
pub final_config: SpinConfig,
pub final_energy: f64,
pub n_iterations: usize,
pub converged: bool,
pub energy_history: Vec<f64>,
}
pub struct MagneticStructureOptimizer {
pub initial: SpinConfig,
pub energy: EnergyFunctional,
pub optimizer_kind: OptimizerKind,
pub lr: f64,
pub max_iter: usize,
pub tol: f64,
}
impl MagneticStructureOptimizer {
pub fn new(initial: SpinConfig, energy: EnergyFunctional) -> Self {
Self {
initial,
energy,
optimizer_kind: OptimizerKind::Adam,
lr: 0.05,
max_iter: 500,
tol: 1e-12,
}
}
pub fn with_optimizer(mut self, kind: OptimizerKind, lr: f64) -> Self {
self.optimizer_kind = kind;
self.lr = lr;
self
}
pub fn with_max_iter(mut self, n: usize) -> Self {
self.max_iter = n;
self
}
pub fn with_tol(mut self, tol: f64) -> Self {
self.tol = tol;
self
}
pub fn optimize(&mut self) -> Result<StructureOptResult> {
let n = self.initial.n_spins();
let mut params = self.initial.params_flat();
let n_p = params.len();
let mut opt: Box<dyn Optimizer> = match self.optimizer_kind {
OptimizerKind::Sgd { momentum } => Box::new(
Sgd::new(self.lr, momentum)
.map_err(|_| invalid_param("sgd", "invalid lr/momentum"))?,
),
OptimizerKind::Adam => {
let mut a = Adam::default_params(n_p);
a.lr = self.lr;
Box::new(a)
},
OptimizerKind::LBfgs { history } => Box::new(
LBfgs::new(self.lr, history)
.map_err(|_| invalid_param("lbfgs", "invalid lr/history"))?,
),
};
let mut energy_history = Vec::with_capacity(self.max_iter);
let mut prev_e = f64::INFINITY;
let mut converged = false;
for _ in 0..self.max_iter {
let tape = Tape::new();
let leaves: Vec<Var<'_>> = params.iter().map(|&p| Var::leaf(&tape, p)).collect();
let energy_var = self.energy.energy(&tape, &leaves)?;
let energy_val = energy_var.value();
tape.backward(energy_var);
let grads: Vec<f64> = leaves.iter().map(|l| l.grad()).collect();
energy_history.push(energy_val);
opt.step(&mut params, &grads);
if (energy_val - prev_e).abs() < self.tol {
converged = true;
break;
}
prev_e = energy_val;
}
let mut final_config = SpinConfig::new(n);
final_config.set_params(¶ms)?;
let final_energy = *energy_history.last().unwrap_or(&f64::NAN);
let n_iterations = energy_history.len();
Ok(StructureOptResult {
final_config,
final_energy,
n_iterations,
converged,
energy_history,
})
}
}
pub fn find_fm_ground_state(
n_spins: usize,
exchange_j: f64,
h_ext_z: f64,
ms: f64,
) -> Result<SpinConfig> {
if n_spins == 0 {
return Err(invalid_param("n_spins", "must be positive"));
}
let initial = SpinConfig::random(n_spins, 0xABCD_1234);
let energy =
EnergyFunctional::new(exchange_j, 0.0, [0.0, 0.0, h_ext_z], ms).with_chain_bonds(n_spins);
let mut opt = MagneticStructureOptimizer::new(initial, energy)
.with_optimizer(OptimizerKind::Adam, 0.05)
.with_max_iter(800)
.with_tol(1e-14);
let res = opt.optimize()?;
Ok(res.final_config)
}
pub fn find_afm_ground_state(n_spins: usize, exchange_j_afm: f64) -> Result<SpinConfig> {
if n_spins == 0 {
return Err(invalid_param("n_spins", "must be positive"));
}
if exchange_j_afm <= 0.0 {
return Err(invalid_param(
"exchange_j_afm",
"magnitude must be positive",
));
}
let mut initial = SpinConfig::new(n_spins);
for i in 0..n_spins {
initial.theta[i] = if i % 2 == 0 {
0.0
} else {
std::f64::consts::PI
};
initial.phi[i] = 0.0;
}
let energy =
EnergyFunctional::new(-exchange_j_afm, 0.0, [0.0, 0.0, 0.0], 0.0).with_chain_bonds(n_spins);
let mut opt = MagneticStructureOptimizer::new(initial, energy)
.with_optimizer(OptimizerKind::Adam, 0.02)
.with_max_iter(400)
.with_tol(1e-14);
let res = opt.optimize()?;
Ok(res.final_config)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::autodiff::tape::Tape;
#[test]
fn test_spin_config_roundtrip_vectors() {
let cfg = SpinConfig::random(8, 99);
let vs = cfg.to_vectors();
let cfg2 = SpinConfig::from_vectors(&vs).unwrap();
let vs2 = cfg2.to_vectors();
for (v1, v2) in vs.iter().zip(vs2.iter()).take(cfg.n_spins()) {
assert!((v1.x - v2.x).abs() < 1e-10);
assert!((v1.y - v2.y).abs() < 1e-10);
assert!((v1.z - v2.z).abs() < 1e-10);
}
}
#[test]
fn test_from_vectors_zero_rejected() {
let v = vec![Vector3::new(0.0_f64, 0.0, 0.0)];
assert!(SpinConfig::from_vectors(&v).is_err());
}
#[test]
fn test_heisenberg_sign() {
let cfg_parallel = SpinConfig::new(2);
let mut cfg_anti = SpinConfig::new(2);
cfg_anti.theta[1] = std::f64::consts::PI;
let ef = EnergyFunctional::new(1.0, 0.0, [0.0, 0.0, 0.0], 0.0).with_chain_bonds(2);
let e_para = ef.energy_f64(&cfg_parallel);
let e_anti = ef.energy_f64(&cfg_anti);
assert!((e_para - (-1.0)).abs() < 1e-12);
assert!((e_anti - 1.0).abs() < 1e-12);
}
#[test]
fn test_energy_gradient_finite_diff() {
let cfg = SpinConfig::random(3, 17);
let ef = EnergyFunctional::new(0.5, 0.1, [0.0, 0.0, 1e3], 1e5).with_chain_bonds(3);
let params = cfg.params_flat();
let tape = Tape::new();
let leaves: Vec<Var<'_>> = params.iter().map(|&p| Var::leaf(&tape, p)).collect();
let e_var = ef.energy(&tape, &leaves).unwrap();
tape.backward(e_var);
let ad: Vec<f64> = leaves.iter().map(|l| l.grad()).collect();
let h = 1e-6;
for k in 0..params.len() {
let mut p_plus = params.clone();
let mut p_minus = params.clone();
p_plus[k] += h;
p_minus[k] -= h;
let mut cfg_p = cfg.clone();
let mut cfg_m = cfg.clone();
cfg_p.set_params(&p_plus).unwrap();
cfg_m.set_params(&p_minus).unwrap();
let fd = (ef.energy_f64(&cfg_p) - ef.energy_f64(&cfg_m)) / (2.0 * h);
assert!(
(ad[k] - fd).abs() < 1e-4,
"k={} AD {} vs FD {}",
k,
ad[k],
fd
);
}
}
#[test]
fn test_fm_ground_state_parallel() {
let cfg = find_fm_ground_state(4, 1.0, 1e6, 1e5).unwrap();
let vs = cfg.to_vectors();
for i in 1..vs.len() {
let dot = vs[0].dot(&vs[i]);
assert!(
dot > 0.95,
"spin {} not parallel to spin 0 (dot = {})",
i,
dot
);
}
}
#[test]
fn test_afm_ground_state_alternates() {
let cfg = find_afm_ground_state(4, 1.0).unwrap();
let vs = cfg.to_vectors();
for i in 0..vs.len() - 1 {
let dot = vs[i].dot(&vs[i + 1]);
assert!(
dot < -0.95,
"bond ({}, {}) not antiparallel (dot = {})",
i,
i + 1,
dot
);
}
}
#[test]
fn test_uniaxial_anisotropy_aligns_z() {
let mut initial = SpinConfig::new(1);
initial.theta[0] = std::f64::consts::FRAC_PI_2 - 0.1;
let ef = EnergyFunctional::new(0.0, 1.0, [0.0, 0.0, 0.0], 0.0);
let mut opt = MagneticStructureOptimizer::new(initial, ef)
.with_optimizer(OptimizerKind::Adam, 0.05)
.with_max_iter(2000)
.with_tol(1e-15);
let res = opt.optimize().unwrap();
let v = res.final_config.to_vectors()[0];
assert!(v.z.abs() > 0.95, "|m_z| = {} not near 1", v.z.abs());
}
#[test]
fn test_zeeman_minimised_when_parallel() {
let mut initial = SpinConfig::new(1);
initial.theta[0] = 0.0; let mut initial_anti = SpinConfig::new(1);
initial_anti.theta[0] = std::f64::consts::PI; let ef = EnergyFunctional::new(0.0, 0.0, [0.0, 0.0, 1e6], 1e5);
let e_para = ef.energy_f64(&initial);
let e_anti = ef.energy_f64(&initial_anti);
assert!(e_para < e_anti);
}
#[test]
fn test_ring_topology_bond_count() {
let ef = EnergyFunctional::new(1.0, 0.0, [0.0, 0.0, 0.0], 0.0).with_ring_bonds(5);
assert_eq!(ef.bonds.len(), 5);
let chain = EnergyFunctional::new(1.0, 0.0, [0.0, 0.0, 0.0], 0.0).with_chain_bonds(5);
assert_eq!(chain.bonds.len(), 4);
}
#[test]
fn test_convergence_within_tol() {
let mut initial = SpinConfig::new(1);
initial.theta[0] = 0.0;
let ef = EnergyFunctional::new(0.0, 1.0, [0.0, 0.0, 1e3], 1e3);
let mut opt = MagneticStructureOptimizer::new(initial, ef)
.with_optimizer(OptimizerKind::Adam, 0.05)
.with_max_iter(1000)
.with_tol(1e-10);
let res = opt.optimize().unwrap();
assert!(res.converged);
assert!(res.n_iterations < 1000);
}
#[test]
fn test_spin_config_params_roundtrip() {
let cfg = SpinConfig::random(5, 31);
let flat = cfg.params_flat();
let mut cfg2 = SpinConfig::new(5);
cfg2.set_params(&flat).unwrap();
for i in 0..5 {
assert!((cfg.theta[i] - cfg2.theta[i]).abs() < 1e-15);
assert!((cfg.phi[i] - cfg2.phi[i]).abs() < 1e-15);
}
}
}