use algorithms_fixtures::dense_fixtures::{heisenberg_mpo_f64, random_mps_unknown_f64};
use ariadnetor_algorithms::dmrg::{
DmrgEnvs, DmrgError, DmrgSweepParams, LocalEigensolverParams, dmrg_2site, sweep_2site,
};
use ariadnetor_algorithms::krylov::LanczosParams;
use ariadnetor_linalg::TruncSvdParams;
use ariadnetor_mps::{CanonicalForm, Mpo, Mps, TensorChain};
use ariadnetor_native::NativeBackend;
use ariadnetor_tensor::{DenseLayout, DenseStorage};
fn small_params(seed: u64) -> DmrgSweepParams {
DmrgSweepParams {
max_sweeps: 4,
min_sweeps: 1,
energy_tol: 1e-10,
eigensolver: LocalEigensolverParams::Lanczos(LanczosParams {
max_iter: 80,
tol: 1e-10,
seed: Some(seed),
}),
trunc: TruncSvdParams {
chi_max: Some(16),
target_trunc_err: None,
},
}
}
#[test]
fn wrapper_dense_heisenberg_n4_matches_manual() {
let n = 4;
let mpo = heisenberg_mpo_f64(n, 1.0);
let psi0 = random_mps_unknown_f64(n, 2, 4, 0xC4F1);
let params = small_params(0xACED);
let mut psi_manual = psi0.clone();
psi_manual.canonicalize(&NativeBackend::new(), 0);
let mut envs_manual = DmrgEnvs::build(&psi_manual, &mpo).expect("manual envs build");
let result_manual =
sweep_2site(&mut envs_manual, &mut psi_manual, &mpo, ¶ms).expect("manual sweep");
let (result_wrapper, psi_wrapper) = dmrg_2site(&mpo, &psi0, ¶ms).expect("wrapper");
assert_eq!(result_wrapper.energy, result_manual.energy);
assert_eq!(result_wrapper.n_sweeps, result_manual.n_sweeps);
assert_eq!(result_wrapper.converged, result_manual.converged);
assert_eq!(
*psi_wrapper.canonical_form(),
CanonicalForm::Mixed { center: 0 },
);
assert_eq!(result_wrapper.sweeps.len(), result_manual.sweeps.len());
for (sw_w, sw_m) in result_wrapper
.sweeps
.iter()
.zip(result_manual.sweeps.iter())
{
assert_eq!(sw_w.sweep_energy, sw_m.sweep_energy);
assert_eq!(sw_w.max_bond, sw_m.max_bond);
assert_eq!(sw_w.steps.len(), sw_m.steps.len());
for (st_w, st_m) in sw_w.steps.iter().zip(sw_m.steps.iter()) {
assert_eq!(st_w.eigenvalue, st_m.eigenvalue);
assert_eq!(st_w.bond_dim, st_m.bond_dim);
assert_eq!(st_w.eigensolver_iters, st_m.eigensolver_iters);
assert_eq!(st_w.eigensolver_converged, st_m.eigensolver_converged);
}
}
for i in 0..n {
assert_eq!(
psi_wrapper.site(i).data_slice(),
psi_manual.site(i).data_slice(),
"site {i} tensor data diverged between wrapper and manual",
);
}
}
#[test]
fn wrapper_accepts_unknown_canonical() {
let n = 4;
let mpo = heisenberg_mpo_f64(n, 1.0);
let psi0 = random_mps_unknown_f64(n, 2, 4, 0xC4F2);
assert_eq!(*psi0.canonical_form(), CanonicalForm::Unknown);
let params = small_params(0xACED);
let (_result, psi_out) =
dmrg_2site(&mpo, &psi0, ¶ms).expect("wrapper should accept Unknown canonical form");
assert_eq!(
*psi_out.canonical_form(),
CanonicalForm::Mixed { center: 0 },
);
}
#[test]
fn wrapper_rejects_empty_mps() {
let mpo: Mpo<DenseStorage<f64>, DenseLayout> = Mpo::empty();
let psi0: Mps<DenseStorage<f64>, DenseLayout> = Mps::empty();
let params = small_params(0xACED);
match dmrg_2site(&mpo, &psi0, ¶ms) {
Err(DmrgError::EmptyMps) => {}
Err(e) => panic!("expected DmrgError::EmptyMps, got Err({e:?})"),
Ok(_) => panic!("expected DmrgError::EmptyMps, got Ok(_)"),
}
}
#[test]
fn wrapper_rejects_length_mismatch() {
let mpo = heisenberg_mpo_f64(4, 1.0);
let psi0 = random_mps_unknown_f64(3, 2, 4, 0xC4F3);
let params = small_params(0xACED);
match dmrg_2site(&mpo, &psi0, ¶ms) {
Err(DmrgError::LengthMismatch { mps: 3, mpo: 4 }) => {}
Err(e) => panic!("expected LengthMismatch {{ mps: 3, mpo: 4 }}, got Err({e:?})"),
Ok(_) => panic!("expected LengthMismatch {{ mps: 3, mpo: 4 }}, got Ok(_)"),
}
}
#[test]
fn wrapper_does_not_mutate_input() {
let n = 4;
let mpo = heisenberg_mpo_f64(n, 1.0);
let psi0 = random_mps_unknown_f64(n, 2, 4, 0xC4F4);
let psi0_snapshot = psi0.clone();
let params = small_params(0xACED);
let _ = dmrg_2site(&mpo, &psi0, ¶ms).expect("wrapper");
assert_eq!(*psi0.canonical_form(), *psi0_snapshot.canonical_form());
for i in 0..n {
assert_eq!(
psi0.site(i).data_slice(),
psi0_snapshot.site(i).data_slice(),
"site {i} tensor data was mutated by the wrapper",
);
}
}