use crate::error::{Result, TbError};
use crate::model::NoRMatrix;
use crate::model_utils::find_R;
use crate::ndarray_lapack::eigvalsh_v;
use crate::{Gauge, Model, OrbitalId, RMatrixData};
use ndarray::parallel::prelude::IntoParallelIterator;
use ndarray::prelude::*;
use ndarray::*;
use ndarray_linalg::UPLO;
use num_complex::Complex;
use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
use std::f64::consts::TAU;
use std::sync::atomic::{AtomicBool, Ordering};
#[derive(Clone, Debug)]
pub struct LightMode {
pub harmonic: isize,
pub a_complex: Array1<Complex<f64>>,
}
impl LightMode {
pub fn new(harmonic: isize, a_complex: Array1<Complex<f64>>) -> Self {
Self {
harmonic,
a_complex,
}
}
}
#[derive(Clone, Debug)]
pub struct FloquetDrive {
pub omega0_ev: f64,
pub modes: Vec<LightMode>,
}
impl FloquetDrive {
pub fn new(omega0_ev: f64) -> Self {
Self {
omega0_ev,
modes: Vec::new(),
}
}
pub fn with_modes(omega0_ev: f64, modes: Vec<LightMode>) -> Self {
Self { omega0_ev, modes }
}
pub fn add_mode(&mut self, mode: LightMode) {
self.modes.push(mode);
}
}
#[derive(Clone, Copy, Debug)]
pub struct FloquetTruncation {
pub n_max: isize,
pub n_time: usize,
}
impl FloquetTruncation {
pub fn new(n_max: isize, n_time: usize) -> Self {
Self { n_max, n_time }
}
#[inline]
pub fn n_sector(&self) -> usize {
(2 * self.n_max + 1) as usize
}
#[inline]
pub fn sectors(&self) -> impl Iterator<Item = isize> {
-self.n_max..=self.n_max
}
}
#[derive(Clone, Debug)]
pub struct IncidentBasis {
pub k_hat: Array1<f64>,
pub e1: Array1<f64>,
pub e2: Array1<f64>,
}
impl IncidentBasis {
pub fn from_direction(k_hat_cart: &Array1<f64>) -> Result<Self> {
if k_hat_cart.len() != 3 {
return Err(TbError::DimensionMismatch {
context: "IncidentBasis::from_direction".to_string(),
expected: 3,
found: k_hat_cart.len(),
});
}
let k_hat = normalize3(k_hat_cart)?;
let reference = if k_hat[2].abs() < 0.9 {
arr1(&[0.0, 0.0, 1.0])
} else {
arr1(&[1.0, 0.0, 0.0])
};
let e1 = normalize3(&cross3(&reference, &k_hat))?;
let e2 = normalize3(&cross3(&k_hat, &e1))?;
Ok(Self { k_hat, e1, e2 })
}
pub fn polarization(&self, jones: [Complex<f64>; 2]) -> Array1<Complex<f64>> {
let mut out = Array1::<Complex<f64>>::zeros(3);
for i in 0..3 {
out[i] = jones[0] * self.e1[i] + jones[1] * self.e2[i];
}
out
}
}
#[derive(Clone, Debug)]
pub struct FloquetEffectiveOptions {
pub order: usize,
pub q_max: Option<isize>,
pub(crate) target_hamR: Option<Array2<isize>>,
}
impl Default for FloquetEffectiveOptions {
fn default() -> Self {
Self {
order: 1,
q_max: None,
target_hamR: None,
}
}
}
impl FloquetEffectiveOptions {
pub fn new() -> Self {
Self::default()
}
pub fn with_order(mut self, order: usize) -> Self {
self.order = order;
self
}
pub fn with_q_max(mut self, q_max: isize) -> Self {
self.q_max = Some(q_max);
self
}
#[cfg(test)]
pub(crate) fn with_target_hamR(mut self, target_hamR: Array2<isize>) -> Self {
self.target_hamR = Some(target_hamR);
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum PeierlsFourierMethod {
TimeGrid,
Bessel {
cutoff_margin: isize,
},
}
struct FloquetHarmonicCache {
q_min: isize,
q_max: isize,
blocks: Array4<Complex<f64>>,
}
impl FloquetHarmonicCache {
#[inline]
fn q_index(&self, q: isize) -> usize {
debug_assert!(
q >= self.q_min && q <= self.q_max,
"Floquet harmonic q={q} is outside cached range [{}, {}]",
self.q_min,
self.q_max
);
(q - self.q_min) as usize
}
}
struct FloquetTimeGrid {
link_field: Array2<f64>,
fourier: Array2<Complex<f64>>,
inv_n_time: f64,
}
impl FloquetTimeGrid {
fn new(
drive: &FloquetDrive,
trunc: &FloquetTruncation,
q_min: isize,
q_max: isize,
dim: usize,
) -> Self {
let n_time = trunc.n_time;
let q_count = (q_max - q_min + 1) as usize;
let inv_n_time = 1.0 / (n_time as f64);
let mut link_field = Array2::<f64>::zeros((n_time, dim));
let mut fourier = Array2::<Complex<f64>>::zeros((q_count, n_time));
for it in 0..n_time {
let theta = TAU * (it as f64) * inv_n_time;
for mode in &drive.modes {
let harmonic_phase = Complex::new(0.0, -(mode.harmonic as f64) * theta).exp();
for a in 0..dim {
link_field[[it, a]] += (mode.a_complex[a] * harmonic_phase).re;
}
}
for (iq, q) in (q_min..=q_max).enumerate() {
fourier[[iq, it]] = Complex::new(0.0, (q as f64) * theta).exp();
}
}
Self {
link_field,
fourier,
inv_n_time,
}
}
}
pub trait Floquet {
type FloquetModel;
fn floquet_model(
&self,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
) -> Result<Self::FloquetModel>;
fn floquet_ham_onek<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
gauge: Gauge,
) -> Result<Array2<Complex<f64>>>;
fn floquet_band_onek<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
gauge: Gauge,
) -> Result<Array1<f64>>;
fn floquet_quasienergy_onek<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
gauge: Gauge,
) -> Result<Array1<f64>>;
}
impl<const SPIN: bool, const DIM: usize, R: RMatrixData> Floquet for Model<SPIN, DIM, R> {
type FloquetModel = Model<SPIN, DIM, NoRMatrix>;
fn floquet_model(
&self,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
) -> Result<Self::FloquetModel> {
validate_floquet_drive::<DIM>(drive, trunc)?;
let nsta = self.nsta();
let norb = self.norb();
let sectors: Vec<isize> = trunc.sectors().collect();
let n_sector = sectors.len();
let new_norb = norb * n_sector;
let total = nsta * n_sector;
let basis_indices = floquet_basis_indices::<SPIN>(nsta, norb, n_sector);
let q_min = -2 * trunc.n_max;
let q_max = 2 * trunc.n_max;
let harmonic_cache = self.floquet_harmonic_cache(
drive,
trunc,
q_min,
q_max,
&PeierlsFourierMethod::TimeGrid,
);
let mut orb = Array2::<f64>::zeros((new_norb, DIM));
for isec in 0..n_sector {
for iorb in 0..norb {
let out_i = isec * norb + iorb;
orb.row_mut(out_i).assign(&self.orb.row(iorb));
}
}
let mut ham_r = self.hamR.clone();
let mut ham = Array3::<Complex<f64>>::zeros((ham_r.nrows(), total, total));
ham.axis_iter_mut(Axis(0))
.into_par_iter()
.enumerate()
.for_each(|(i_r, mut out)| {
for i in 0..nsta {
for j in 0..nsta {
if harmonic_cache
.blocks
.slice(s![.., i_r, i, j])
.iter()
.all(|x| x.norm_sqr() == 0.0)
{
continue;
}
for (in_sec, &n) in sectors.iter().enumerate() {
let row = basis_indices[in_sec][i];
for (im_sec, &m) in sectors.iter().enumerate() {
let hopping = harmonic_cache.blocks
[[harmonic_cache.q_index(n - m), i_r, i, j]];
if hopping.norm_sqr() == 0.0 {
continue;
}
let col = basis_indices[im_sec][j];
out[[row, col]] += hopping;
}
}
}
}
});
let zero_r = Array1::<isize>::zeros(DIM);
let onsite_index = match find_R(&ham_r, &zero_r) {
Some(index) => index,
None => {
ham.push(
Axis(0),
Array2::<Complex<f64>>::zeros((total, total)).view(),
)
.unwrap();
ham_r.push_row(zero_r.view()).unwrap();
ham.len_of(Axis(0)) - 1
}
};
for (in_sec, &n) in sectors.iter().enumerate() {
let photon_shift = n as f64 * drive.omega0_ev;
for i in 0..nsta {
let idx = basis_indices[in_sec][i];
ham[[onsite_index, idx, idx]] += Complex::new(photon_shift, 0.0);
}
}
let atoms = (0..n_sector)
.flat_map(|sector| {
self.atoms.iter().cloned().map(move |mut atom| {
atom.set_orbitals(
atom.orbitals()
.iter()
.map(|id| OrbitalId::new(sector * norb + id.index()))
.collect(),
);
atom
})
})
.collect();
let mut model =
Model::<SPIN, DIM, NoRMatrix>::tb_model(self.lat.clone(), orb, Some(atoms))?;
model.ham = ham;
model.hamR = ham_r;
model.orb_projection = (0..n_sector)
.flat_map(|_| (0..norb).map(|i| self.orb_projection[i]))
.collect();
Ok(model)
}
fn floquet_ham_onek<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
gauge: Gauge,
) -> Result<Array2<Complex<f64>>> {
validate_floquet_input(self, kvec, drive, trunc)?;
let nsta = self.nsta();
let norb = self.norb();
let n_sector = trunc.n_sector();
let total = nsta * n_sector;
let mut hamf = Array2::<Complex<f64>>::zeros((total, total));
let basis_indices = floquet_basis_indices::<SPIN>(nsta, norb, n_sector);
let q_min = -2 * trunc.n_max;
let q_max = 2 * trunc.n_max;
let harmonic_cache = self.floquet_harmonic_cache(
drive,
trunc,
q_min,
q_max,
&PeierlsFourierMethod::TimeGrid,
);
let hq: Vec<Array2<Complex<f64>>> = (q_min..=q_max)
.map(|q| self.floquet_cached_harmonic_onek(kvec, q, gauge, &harmonic_cache))
.collect();
for (in_sec, n) in trunc.sectors().enumerate() {
for (im_sec, m) in trunc.sectors().enumerate() {
let q = n - m;
let block = &hq[(q - q_min) as usize];
for i in 0..nsta {
for j in 0..nsta {
let row = basis_indices[in_sec][i];
let col = basis_indices[im_sec][j];
hamf[[row, col]] = block[[i, j]];
}
}
}
let photon_shift = n as f64 * drive.omega0_ev;
for i in 0..nsta {
let idx = basis_indices[in_sec][i];
hamf[[idx, idx]] += photon_shift;
}
}
Ok(hamf)
}
fn floquet_band_onek<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
gauge: Gauge,
) -> Result<Array1<f64>> {
let hamf = self.floquet_ham_onek(kvec, drive, trunc, gauge)?;
Ok(eigvalsh_v(&hamf, UPLO::Upper))
}
fn floquet_quasienergy_onek<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
gauge: Gauge,
) -> Result<Array1<f64>> {
let mut values = self.floquet_band_onek(kvec, drive, trunc, gauge)?;
values.mapv_inplace(|x| fold_quasienergy(x, drive.omega0_ev));
values
.as_slice_mut()
.unwrap()
.sort_by(|a, b| a.partial_cmp(b).unwrap());
Ok(values)
}
}
impl<const SPIN: bool, const DIM: usize, R: RMatrixData> Model<SPIN, DIM, R> {
#[cfg(test)]
pub(crate) fn floquet_effective_model_legacy(
&self,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
k_mesh: [usize; DIM],
options: Option<&FloquetEffectiveOptions>,
) -> Result<Model<SPIN, DIM, NoRMatrix>> {
let default_options;
let options = match options {
Some(options) => options,
None => {
default_options = FloquetEffectiveOptions::default();
&default_options
}
};
validate_floquet_drive::<DIM>(drive, trunc)?;
validate_effective_options::<DIM>(&k_mesh, options)?;
let nsta = self.nsta();
let target_ham_r = options
.target_hamR
.clone()
.unwrap_or_else(|| self.hamR.clone());
validate_target_hamr::<DIM>(&target_ham_r)?;
let q_max = options.q_max.unwrap_or(2 * trunc.n_max);
if q_max < 0 {
return Err(TbError::Other(format!(
"FloquetEffectiveOptions.q_max must be non-negative, got {q_max}"
)));
}
let harmonic_cache = self.floquet_harmonic_cache(
drive,
trunc,
-q_max,
q_max,
&PeierlsFourierMethod::TimeGrid,
);
let kpoints = floquet_uniform_kmesh(&k_mesh);
let norm = 1.0 / (kpoints.len() as f64);
let ham = kpoints
.par_iter()
.fold(
|| Array3::<Complex<f64>>::zeros((target_ham_r.nrows(), nsta, nsta)),
|mut partial, kvec| {
let h_eff = self.floquet_effective_ham_onek_lattice(
kvec,
drive,
options.order,
q_max,
&harmonic_cache,
);
for (i_r, r_vec) in target_ham_r.outer_iter().enumerate() {
let phase = inverse_bloch_phase::<DIM, _>(&r_vec, kvec) * norm;
let mut block = partial.index_axis_mut(Axis(0), i_r);
crate::ndarray_lapack::zaxpy(
phase,
h_eff.as_slice().unwrap(),
block.as_slice_mut().unwrap(),
);
}
partial
},
)
.reduce(
|| Array3::<Complex<f64>>::zeros((target_ham_r.nrows(), nsta, nsta)),
|mut left, right| {
left.zip_mut_with(&right, |a, b| *a += *b);
left
},
);
let mut ham = ham;
enforce_real_space_hermiticity(&mut ham, &target_ham_r)?;
let mut model = Model::<SPIN, DIM, NoRMatrix>::tb_model(
self.lat.clone(),
self.orb.clone(),
Some(self.atoms.clone()),
)?;
model.ham = ham;
model.hamR = target_ham_r;
model.orb_projection = self.orb_projection.clone();
Ok(model)
}
pub fn floquet_effective_model(
&self,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
options: Option<&FloquetEffectiveOptions>,
) -> Result<Model<SPIN, DIM, NoRMatrix>> {
let default_options;
let options = match options {
Some(options) => options,
None => {
default_options = FloquetEffectiveOptions::default();
&default_options
}
};
validate_floquet_drive::<DIM>(drive, trunc)?;
if options.order > 1 {
return Err(TbError::Other(format!(
"FloquetEffectiveOptions.order must be 0 or 1, got {}",
options.order
)));
}
if let Some(target) = &options.target_hamR {
return Err(TbError::Other(format!(
"FloquetEffectiveOptions.target_hamR is not supported by the \
real-space path: the effective support is determined \
automatically (got {} target vectors)",
target.nrows()
)));
}
let q_max = options.q_max.unwrap_or(2 * trunc.n_max);
if q_max < 0 {
return Err(TbError::Other(format!(
"FloquetEffectiveOptions.q_max must be non-negative, got {q_max}"
)));
}
let nsta = self.nsta();
let harmonic_cache = self.floquet_harmonic_cache(
drive,
trunc,
-q_max,
q_max,
&PeierlsFourierMethod::Bessel { cutoff_margin: 6 },
);
let mut blocks = std::collections::BTreeMap::<Vec<isize>, Array2<Complex<f64>>>::new();
let iq0 = harmonic_cache.q_index(0);
for (i_r, row) in self.hamR.outer_iter().enumerate() {
blocks.insert(
row.to_vec(),
harmonic_cache.blocks.slice(s![iq0, i_r, .., ..]).to_owned(),
);
}
if options.order == 1 {
for q in 1..=q_max {
let q_idx = harmonic_cache.q_index(q);
let m_idx = harmonic_cache.q_index(-q);
let a_blocks: Vec<Array2<Complex<f64>>> = (0..self.hamR.nrows())
.map(|i_r| {
harmonic_cache
.blocks
.slice(s![q_idx, i_r, .., ..])
.to_owned()
})
.collect();
let b_blocks: Vec<Array2<Complex<f64>>> = (0..self.hamR.nrows())
.map(|i_r| {
harmonic_cache
.blocks
.slice(s![m_idx, i_r, .., ..])
.to_owned()
})
.collect();
let (comm_blocks, comm_r) =
real_space_commutator(&a_blocks, &b_blocks, &self.hamR)?;
let scale = 1.0 / ((q as f64) * drive.omega0_ev);
for (i_r, row) in comm_r.outer_iter().enumerate() {
let contribution = comm_blocks[i_r].mapv(|x| x * scale);
blocks
.entry(row.to_vec())
.and_modify(|block| *block += &contribution)
.or_insert(contribution);
}
}
}
let n_r_out = blocks.len();
let mut ham = Array3::<Complex<f64>>::zeros((n_r_out, nsta, nsta));
let mut ham_r = Array2::<isize>::zeros((n_r_out, DIM));
for (i, (key, block)) in blocks.into_iter().enumerate() {
for (a, v) in key.iter().enumerate() {
ham_r[[i, a]] = *v;
}
ham.index_axis_mut(Axis(0), i).assign(&block);
}
enforce_real_space_hermiticity(&mut ham, &ham_r)?;
let mut model = Model::<SPIN, DIM, NoRMatrix>::tb_model(
self.lat.clone(),
self.orb.clone(),
Some(self.atoms.clone()),
)?;
model.ham = ham;
model.hamR = ham_r;
model.orb_projection = self.orb_projection.clone();
Ok(model)
}
#[cfg(test)]
fn floquet_effective_ham_onek_lattice<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
order: usize,
q_max: isize,
harmonic_cache: &FloquetHarmonicCache,
) -> Array2<Complex<f64>> {
let mut h_eff = self.floquet_cached_harmonic_onek(kvec, 0, Gauge::Lattice, harmonic_cache);
match order {
0 => {}
1 => {
for q in 1..=q_max {
let h_pos =
self.floquet_cached_harmonic_onek(kvec, q, Gauge::Lattice, harmonic_cache);
let h_neg =
self.floquet_cached_harmonic_onek(kvec, -q, Gauge::Lattice, harmonic_cache);
let comm = h_pos.dot(&h_neg) - h_neg.dot(&h_pos);
h_eff = h_eff + comm.mapv(|x| x / ((q as f64) * drive.omega0_ev));
}
}
_ => unreachable!("effective order is validated before evaluation"),
}
h_eff
}
fn floquet_harmonic_cache(
&self,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
q_min: isize,
q_max: isize,
method: &PeierlsFourierMethod,
) -> FloquetHarmonicCache {
let nsta = self.nsta();
let norb = self.norb();
let n_r = self.hamR.nrows();
let q_count = (q_max - q_min + 1) as usize;
let mut blocks = Array4::<Complex<f64>>::zeros((q_count, n_r, nsta, nsta));
if drive.modes.is_empty() {
if q_min <= 0 && 0 <= q_max {
blocks
.slice_mut(s![(0 - q_min) as usize, .., .., ..])
.assign(&self.ham);
}
return FloquetHarmonicCache {
q_min,
q_max,
blocks,
};
}
let mut d_index = std::collections::HashMap::<[u64; DIM], usize>::new();
let mut unique_d = Vec::<Array1<f64>>::new();
let mut entries = Vec::<(usize, usize, usize, usize)>::new(); for i_r in 0..n_r {
let r_vec = self.hamR.row(i_r);
for i in 0..nsta {
for j in 0..nsta {
if self.ham[[i_r, i, j]].norm_sqr() == 0.0 {
continue;
}
let d_cart = self.link_displacement_cartesian(i % norb, j % norb, &r_vec);
let mut key = [0_u64; DIM];
for a in 0..DIM {
key[a] = d_cart[a].to_bits();
}
let index = *d_index.entry(key).or_insert_with(|| {
unique_d.push(d_cart);
unique_d.len() - 1
});
entries.push((i_r, i, j, index));
}
}
}
let time_grid = FloquetTimeGrid::new(drive, trunc, q_min, q_max, DIM);
let fallback_warned = AtomicBool::new(false);
let fallback_clamped = AtomicBool::new(false);
let coeffs_per_d: Vec<Array1<Complex<f64>>> = unique_d
.par_iter()
.map(|d| match method {
PeierlsFourierMethod::Bessel { cutoff_margin } => {
match bessel_peierls_coeffs(d, drive, q_min, q_max, *cutoff_margin) {
Ok(coeffs) => coeffs,
Err(error) => {
if !fallback_warned.swap(true, Ordering::Relaxed) {
eprintln!(
"Bessel backend unavailable for some links \
({error}); falling back to the time grid"
);
}
fallback_time_grid_coeffs(
d,
drive,
trunc,
q_min,
q_max,
DIM,
&time_grid,
&fallback_clamped,
)
}
}
}
PeierlsFourierMethod::TimeGrid => {
Array1::from(peierls_fourier_coeffs(d, q_min, q_max, drive, &time_grid))
}
})
.collect();
for (i_r, i, j, d_index) in entries {
let t = self.ham[[i_r, i, j]];
for (iq, coeff) in coeffs_per_d[d_index].iter().enumerate() {
if coeff.norm_sqr() != 0.0 {
blocks[[iq, i_r, i, j]] = t * coeff;
}
}
}
FloquetHarmonicCache {
q_min,
q_max,
blocks,
}
}
fn floquet_cached_harmonic_onek<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
q: isize,
gauge: Gauge,
harmonic_cache: &FloquetHarmonicCache,
) -> Array2<Complex<f64>> {
let nsta = self.nsta();
let mut hamq = Array2::<Complex<f64>>::zeros((nsta, nsta));
let iq = harmonic_cache.q_index(q);
let hamq_slice = hamq.as_slice_mut().unwrap();
for i_r in 0..self.hamR.nrows() {
let r_vec = self.hamR.row(i_r);
let bloch = bloch_phase::<DIM, S>(&r_vec, kvec);
let block = harmonic_cache.blocks.slice(s![iq, i_r, .., ..]);
crate::ndarray_lapack::zaxpy(bloch, block.as_slice().unwrap(), hamq_slice);
}
match gauge {
Gauge::Lattice => hamq,
Gauge::Atom => self.apply_atom_gauge(kvec, hamq),
}
}
fn link_displacement_cartesian(
&self,
i_orb: usize,
j_orb: usize,
r_vec: &ArrayView1<'_, isize>,
) -> Array1<f64> {
let mut frac = Array1::<f64>::zeros(DIM);
for a in 0..DIM {
frac[a] = r_vec[a] as f64 + self.orb[[j_orb, a]] - self.orb[[i_orb, a]];
}
frac.dot(&self.lat)
}
fn apply_atom_gauge<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
mut ham: Array2<Complex<f64>>,
) -> Array2<Complex<f64>> {
let nsta = self.nsta();
let norb = self.norb();
let mut phase_orb = Array1::<Complex<f64>>::zeros(norb);
for i in 0..norb {
let mut tau_dot_k = 0.0;
for a in 0..DIM {
tau_dot_k += self.orb[[i, a]] * kvec[a];
}
phase_orb[i] = Complex::new(0.0, TAU * tau_dot_k).exp();
}
let mut phase = Array1::<Complex<f64>>::zeros(nsta);
phase.slice_mut(s![..norb]).assign(&phase_orb);
if SPIN {
phase.slice_mut(s![norb..]).assign(&phase_orb);
}
for i in 0..nsta {
let left = phase[i].conj();
for j in 0..nsta {
ham[[i, j]] *= left * phase[j];
}
}
ham
}
}
#[inline]
fn floquet_basis_indices<const SPIN: bool>(
nsta: usize,
norb: usize,
n_sector: usize,
) -> Vec<Vec<usize>> {
(0..n_sector)
.map(|sector_index| {
(0..nsta)
.map(|state_index| {
floquet_basis_index::<SPIN>(sector_index, state_index, nsta, norb, n_sector)
})
.collect()
})
.collect()
}
#[inline]
fn floquet_basis_index<const SPIN: bool>(
sector_index: usize,
state_index: usize,
nsta: usize,
norb: usize,
n_sector: usize,
) -> usize {
if SPIN {
let spin = state_index / norb;
let orbital = state_index % norb;
spin * n_sector * norb + sector_index * norb + orbital
} else {
sector_index * nsta + state_index
}
}
#[inline]
pub fn fold_quasienergy(energy: f64, omega0_ev: f64) -> f64 {
(energy + 0.5 * omega0_ev).rem_euclid(omega0_ev) - 0.5 * omega0_ev
}
fn validate_floquet_input<
const DIM: usize,
S: Data<Elem = f64>,
R: RMatrixData,
const SPIN: bool,
>(
model: &Model<SPIN, DIM, R>,
kvec: &ArrayBase<S, Ix1>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
) -> Result<()> {
if kvec.len() != DIM {
return Err(TbError::KVectorLengthMismatch {
expected: DIM,
actual: kvec.len(),
});
}
if model.lat.nrows() != DIM || model.lat.ncols() != DIM {
return Err(TbError::InvalidArrayShape {
expected: vec![DIM, DIM],
found: vec![model.lat.nrows(), model.lat.ncols()],
});
}
validate_floquet_drive::<DIM>(drive, trunc)
}
fn validate_floquet_drive<const DIM: usize>(
drive: &FloquetDrive,
trunc: &FloquetTruncation,
) -> Result<()> {
if !drive.omega0_ev.is_finite() || drive.omega0_ev <= 0.0 {
return Err(TbError::InvalidEnergyRange {
min: 0.0,
max: drive.omega0_ev,
});
}
if trunc.n_max < 0 {
return Err(TbError::Other(format!(
"FloquetTruncation.n_max must be non-negative, got {}",
trunc.n_max
)));
}
if trunc.n_time == 0 {
return Err(TbError::Other(
"FloquetTruncation.n_time must be positive".to_string(),
));
}
for (im, mode) in drive.modes.iter().enumerate() {
if mode.a_complex.len() != DIM {
return Err(TbError::DimensionMismatch {
context: format!("FloquetDrive.modes[{im}].a_complex"),
expected: DIM,
found: mode.a_complex.len(),
});
}
if mode
.a_complex
.iter()
.any(|z| !z.re.is_finite() || !z.im.is_finite())
{
return Err(TbError::Other(format!(
"FloquetDrive.modes[{im}].a_complex contains non-finite values"
)));
}
}
Ok(())
}
#[cfg(test)]
fn validate_effective_options<const DIM: usize>(
k_mesh: &[usize; DIM],
options: &FloquetEffectiveOptions,
) -> Result<()> {
if options.order > 1 {
return Err(TbError::Other(format!(
"Floquet effective order {} is not implemented; supported orders are 0 and 1",
options.order
)));
}
for (axis, &n) in k_mesh.iter().enumerate() {
if n == 0 {
return Err(TbError::Other(format!(
"FloquetEffectiveOptions.k_mesh[{axis}] must be positive"
)));
}
}
if let Some(q_max) = options.q_max {
if q_max < 0 {
return Err(TbError::Other(format!(
"FloquetEffectiveOptions.q_max must be non-negative, got {q_max}"
)));
}
}
if let Some(target_ham_r) = &options.target_hamR {
validate_target_hamr::<DIM>(target_ham_r)?;
}
Ok(())
}
#[cfg(test)]
fn validate_target_hamr<const DIM: usize>(target_ham_r: &Array2<isize>) -> Result<()> {
if target_ham_r.ncols() != DIM {
return Err(TbError::InvalidArrayShape {
expected: vec![target_ham_r.nrows(), DIM],
found: vec![target_ham_r.nrows(), target_ham_r.ncols()],
});
}
if target_ham_r.nrows() == 0 {
return Err(TbError::Other(
"target_hamR must contain at least one R vector".to_string(),
));
}
for i_r in 0..target_ham_r.nrows() {
let r = target_ham_r.row(i_r).to_owned();
if (0..i_r).any(|j_r| {
target_ham_r
.row(j_r)
.iter()
.zip(r.iter())
.all(|(left, right)| left == right)
}) {
return Err(TbError::Other(format!(
"target_hamR contains the duplicate vector R={:?}",
r.to_vec()
)));
}
let neg_r = r.mapv(|x| -x);
if find_R(target_ham_r, &neg_r).is_none() {
return Err(TbError::MissingHermitianConjugateHopping { r });
}
}
Ok(())
}
pub(crate) fn bessel_j(m: isize, r: f64) -> f64 {
assert!(
r.is_finite() && r >= 0.0,
"bessel_j expects a non-negative finite argument, got {r}"
);
if m < 0 {
return if m.rem_euclid(2) == 1 {
-bessel_j(-m, r)
} else {
bessel_j(-m, r)
};
}
if r == 0.0 {
return if m == 0 { 1.0 } else { 0.0 };
}
puruspe::Jn(m as u32, r)
}
fn bessel_adaptive_m_cap(r: f64, error_share: f64, margin: isize) -> isize {
debug_assert!(
r.is_finite() && r >= 0.0,
"bessel_adaptive_m_cap: r must be finite and non-negative"
);
let mut m_cap = (r.ceil() as isize).saturating_add(margin).min(4096);
while m_cap <= 4096 {
let mut tail = 0.0;
let mut current = bessel_j(m_cap + 1, r).abs();
for m in (m_cap + 2)..(m_cap + 201) {
tail += current;
if current < 1e-20 {
break;
}
current = bessel_j(m, r).abs();
}
if 2.0 * tail <= error_share {
return m_cap;
}
m_cap += 1;
}
m_cap.min(4096)
}
pub(crate) fn bessel_peierls_coeffs(
d: &Array1<f64>,
drive: &FloquetDrive,
q_min: isize,
q_max: isize,
cutoff_margin: isize,
) -> Result<Array1<Complex<f64>>> {
if q_min > q_max {
return Err(TbError::Other(format!(
"bessel_peierls_coeffs: empty harmonic range [{q_min}, {q_max}]"
)));
}
if !(0..=48).contains(&cutoff_margin) {
return Err(TbError::Other(format!(
"bessel_peierls_coeffs: cutoff_margin = {cutoff_margin} outside [0, 48]"
)));
}
let q_count = q_max.checked_sub(q_min).ok_or_else(|| {
TbError::Other("bessel_peierls_coeffs: harmonic range too wide".to_string())
})? as usize
+ 1;
if drive.modes.is_empty() {
let mut coeffs = Array1::<Complex<f64>>::zeros(q_count);
if q_min <= 0 && 0 <= q_max {
coeffs[(0 - q_min) as usize] = Complex::new(1.0, 0.0);
}
return Ok(coeffs);
}
struct ModeData {
r: f64,
delta: f64,
harmonic: isize,
m_cap: isize,
}
let mut modes = Vec::<ModeData>::with_capacity(drive.modes.len());
let error_share = 1e-12 / (drive.modes.len() as f64);
let mut total_drift = 0_isize;
for mode in &drive.modes {
let z: Complex<f64> = mode
.a_complex
.iter()
.zip(d.iter())
.map(|(a, d)| *a * *d)
.sum();
let r = z.norm();
if r == 0.0 {
continue;
}
if r > 8.0 {
return Err(TbError::Other(format!(
"bessel_peierls_coeffs: mode amplitude R = {r:.3} exceeds the \
Bessel backend's range (R ≤ 8); use the time-grid backend"
)));
}
let m_cap = bessel_adaptive_m_cap(r, error_share, cutoff_margin);
assert!(
m_cap <= 64,
"Bessel cutoff exceeded the safety cap for R = {r}"
);
let harmonic_abs = mode.harmonic.checked_abs().ok_or_else(|| {
TbError::Other("bessel_peierls_coeffs: harmonic drift overflow".to_string())
})?;
total_drift = total_drift
.checked_add(harmonic_abs.checked_mul(m_cap).ok_or_else(|| {
TbError::Other("bessel_peierls_coeffs: harmonic drift overflow".to_string())
})?)
.ok_or_else(|| {
TbError::Other("bessel_peierls_coeffs: harmonic drift overflow".to_string())
})?;
modes.push(ModeData {
r,
delta: z.arg(),
harmonic: mode.harmonic,
m_cap,
});
}
let work_min = q_min.checked_sub(total_drift).ok_or_else(|| {
TbError::Other("bessel_peierls_coeffs: working window underflow".to_string())
})?;
let work_max = q_max.checked_add(total_drift).ok_or_else(|| {
TbError::Other("bessel_peierls_coeffs: working window overflow".to_string())
})?;
let work_span = work_max.checked_sub(work_min).ok_or_else(|| {
TbError::Other("bessel_peierls_coeffs: working window span overflow".to_string())
})?;
let work_len = usize::try_from(work_span).map_err(|_| {
TbError::Other("bessel_peierls_coeffs: working window too large".to_string())
})? + 1;
let mut sequence = vec![Complex::new(0.0, 0.0); work_len];
if (0_isize..work_len as isize).contains(&(0 - work_min)) {
sequence[(0 - work_min) as usize] = Complex::new(1.0, 0.0);
}
for mode in &modes {
let mut minus_i_power = Complex::new(1.0, 0.0); let mut b = Vec::<(isize, Complex<f64>)>::with_capacity((2 * mode.m_cap + 1) as usize);
for m in 0..=mode.m_cap {
let value = minus_i_power
* bessel_j(m, mode.r)
* Complex::from_polar(1.0, -(m as f64) * mode.delta);
if m == 0 {
b.push((0, value));
} else {
let neg = minus_i_power
* bessel_j(m, mode.r)
* Complex::from_polar(1.0, (m as f64) * mode.delta);
b.push((-m, neg));
b.push((m, value));
}
minus_i_power *= Complex::new(0.0, -1.0); }
let mut next = vec![Complex::new(0.0, 0.0); work_len];
for &(m, weight) in &b {
let shift = mode.harmonic * m;
for (index, _) in sequence.iter().enumerate() {
let q = work_min + index as isize;
let Some(source) = q.checked_add(shift) else {
continue;
};
let Some(source_index) = source.checked_sub(work_min) else {
continue;
};
if source_index >= 0 && (source_index as usize) < work_len {
next[index] += sequence[source_index as usize] * weight;
}
}
}
sequence = next;
}
Ok(Array1::from(
sequence[(q_min - work_min) as usize..(q_max - work_min + 1) as usize].to_vec(),
))
}
fn peierls_fourier_coeffs(
d_cart: &Array1<f64>,
q_min: isize,
q_max: isize,
drive: &FloquetDrive,
time_grid: &FloquetTimeGrid,
) -> Vec<Complex<f64>> {
let q_count = (q_max - q_min + 1) as usize;
if drive.modes.is_empty() {
let mut coeffs = vec![Complex::new(0.0, 0.0); q_count];
if q_min <= 0 && 0 <= q_max {
coeffs[(0 - q_min) as usize] = Complex::new(1.0, 0.0);
}
return coeffs;
}
let mut coeffs = vec![Complex::new(0.0, 0.0); q_count];
for it in 0..time_grid.link_field.nrows() {
let mut link_phase = 0.0;
for a in 0..d_cart.len() {
link_phase += time_grid.link_field[[it, a]] * d_cart[a];
}
let peierls = Complex::new(0.0, -link_phase).exp();
for (iq, coeff) in coeffs.iter_mut().enumerate() {
*coeff += time_grid.fourier[[iq, it]] * peierls;
}
}
for coeff in &mut coeffs {
*coeff *= time_grid.inv_n_time;
}
coeffs
}
fn fallback_time_grid_coeffs(
d: &Array1<f64>,
drive: &FloquetDrive,
trunc: &FloquetTruncation,
q_min: isize,
q_max: isize,
dim: usize,
time_grid: &FloquetTimeGrid,
clamped: &AtomicBool,
) -> Array1<Complex<f64>> {
let mut bandwidth = 0_usize;
for mode in &drive.modes {
let z: Complex<f64> = mode
.a_complex
.iter()
.zip(d.iter())
.map(|(a, d)| *a * *d)
.sum();
let r = z.norm();
if r == 0.0 {
continue;
}
if !r.is_finite() {
continue;
}
let m_cap = bessel_adaptive_m_cap(r, 1e-12, 48);
let drift = (mode.harmonic.unsigned_abs() as usize).saturating_mul(m_cap as usize);
bandwidth = bandwidth.saturating_add(drift);
}
let required = bandwidth.saturating_mul(2).saturating_add(4);
if required <= trunc.n_time {
return Array1::from(peierls_fourier_coeffs(d, q_min, q_max, drive, time_grid));
}
const FALLBACK_GRID_MAX: usize = 1 << 20;
let n_req = required.min(FALLBACK_GRID_MAX);
if required > FALLBACK_GRID_MAX && !clamped.swap(true, Ordering::Relaxed) {
eprintln!(
"Floquet fallback grid clamped to {FALLBACK_GRID_MAX} time points \
(link requires {required}); coefficients on this link may be \
inaccurate"
);
}
let q_count = (q_max - q_min + 1) as usize;
let inv_n = 1.0 / (n_req as f64);
let mut coeffs = vec![Complex::new(0.0, 0.0); q_count];
for it in 0..n_req {
let theta = TAU * (it as f64) * inv_n;
let mut link_phase = 0.0;
for mode in &drive.modes {
let harmonic_phase = Complex::new(0.0, -(mode.harmonic as f64) * theta).exp();
for a in 0..dim {
link_phase += (mode.a_complex[a] * harmonic_phase).re * d[a];
}
}
let peierls = Complex::new(0.0, -link_phase).exp();
for (iq, q) in (q_min..=q_max).enumerate() {
coeffs[iq] += Complex::new(0.0, (q as f64) * theta).exp() * peierls;
}
}
for coeff in &mut coeffs {
*coeff *= inv_n;
}
Array1::from(coeffs)
}
fn zgemm_row_accumulate(
alpha: Complex<f64>,
a: &Array2<Complex<f64>>,
b: &Array2<Complex<f64>>,
c: &mut Array2<Complex<f64>>,
) {
let n = a.nrows();
debug_assert_eq!(
(a.ncols(), b.nrows(), b.ncols(), c.nrows(), c.ncols()),
(n, n, n, n, n),
"zgemm_row_accumulate: square n x n blocks required"
);
let n_i = n as i32;
let beta = Complex::new(1.0, 0.0);
unsafe {
blas::zgemm(
b'N',
b'N',
n_i,
n_i,
n_i,
alpha,
b.as_slice().unwrap(),
n_i,
a.as_slice().unwrap(),
n_i,
beta,
c.as_slice_mut().unwrap(),
n_i,
);
}
}
fn real_space_commutator(
a_blocks: &[Array2<Complex<f64>>],
b_blocks: &[Array2<Complex<f64>>],
ham_r: &Array2<isize>,
) -> Result<(Vec<Array2<Complex<f64>>>, Array2<isize>)> {
let n_r = ham_r.nrows();
debug_assert_eq!(a_blocks.len(), n_r, "a_blocks must match hamR row count");
debug_assert_eq!(b_blocks.len(), n_r, "b_blocks must match hamR row count");
let nsta = a_blocks.first().map_or(0, |block| block.nrows());
for block in a_blocks.iter().chain(b_blocks) {
debug_assert_eq!(
(block.nrows(), block.ncols()),
(nsta, nsta),
"all blocks must be square nsta x nsta"
);
}
let mut support = std::collections::BTreeSet::<Vec<isize>>::new();
for r1 in ham_r.outer_iter() {
for r2 in ham_r.outer_iter() {
support.insert(r1.iter().zip(r2.iter()).map(|(a, b)| a + b).collect());
}
}
let mut support_rows = Array2::<isize>::zeros((support.len(), ham_r.ncols()));
for (i, r) in support.iter().enumerate() {
for (a, v) in r.iter().enumerate() {
support_rows[[i, a]] = *v;
}
}
let mut index = std::collections::HashMap::<Vec<isize>, usize>::with_capacity(n_r);
for (i_r, row) in ham_r.outer_iter().enumerate() {
index.insert(row.to_vec(), i_r);
}
let one = Complex::new(1.0, 0.0);
let minus_one = Complex::new(-1.0, 0.0);
let support_ordered: Vec<Vec<isize>> = support.into_iter().collect();
let blocks: Vec<Array2<Complex<f64>>> = support_ordered
.par_iter()
.map(|r| {
let mut comm = Array2::<Complex<f64>>::zeros((nsta, nsta));
for (i_r2, r2_row) in ham_r.outer_iter().enumerate() {
let r1: Vec<isize> = r.iter().zip(r2_row.iter()).map(|(a, b)| a - b).collect();
let Some(&i_r1) = index.get(&r1) else {
continue;
};
zgemm_row_accumulate(one, &a_blocks[i_r1], &b_blocks[i_r2], &mut comm);
zgemm_row_accumulate(minus_one, &b_blocks[i_r1], &a_blocks[i_r2], &mut comm);
}
comm
})
.collect();
let mut stacked = Array3::<Complex<f64>>::zeros((blocks.len(), nsta, nsta));
for (i, block) in blocks.iter().enumerate() {
stacked.index_axis_mut(Axis(0), i).assign(block);
}
enforce_real_space_hermiticity(&mut stacked, &support_rows)?;
let blocks: Vec<Array2<Complex<f64>>> = (0..blocks.len())
.map(|i| stacked.index_axis(Axis(0), i).to_owned())
.collect();
Ok((blocks, support_rows))
}
fn bloch_phase<const DIM: usize, S: Data<Elem = f64>>(
r_vec: &ArrayView1<'_, isize>,
kvec: &ArrayBase<S, Ix1>,
) -> Complex<f64> {
let mut r_dot_k = 0.0;
for a in 0..DIM {
r_dot_k += r_vec[a] as f64 * kvec[a];
}
Complex::new(0.0, TAU * r_dot_k).exp()
}
#[cfg(test)]
fn inverse_bloch_phase<const DIM: usize, S: Data<Elem = f64>>(
r_vec: &ArrayView1<'_, isize>,
kvec: &ArrayBase<S, Ix1>,
) -> Complex<f64> {
bloch_phase::<DIM, S>(r_vec, kvec).conj()
}
#[cfg(test)]
fn floquet_uniform_kmesh<const DIM: usize>(mesh: &[usize; DIM]) -> Vec<Array1<f64>> {
let n_total = mesh.iter().product();
let mut points = Vec::with_capacity(n_total);
for mut linear in 0..n_total {
let mut k = Array1::<f64>::zeros(DIM);
for a in (0..DIM).rev() {
let n = mesh[a];
let i = linear % n;
linear /= n;
k[a] = (i as f64) / (n as f64);
}
points.push(k);
}
points
}
fn enforce_real_space_hermiticity(
ham: &mut Array3<Complex<f64>>,
ham_r: &Array2<isize>,
) -> Result<()> {
let n_r = ham_r.nrows();
let mut visited = vec![false; n_r];
for i_r in 0..n_r {
if visited[i_r] {
continue;
}
let neg_r = ham_r.row(i_r).mapv(|x| -x);
let Some(j_r) = find_R(ham_r, &neg_r) else {
return Err(TbError::MissingHermitianConjugateHopping {
r: ham_r.row(i_r).to_owned(),
});
};
if i_r == j_r {
let block = ham.index_axis(Axis(0), i_r).to_owned();
let herm = (&block + &hermitian_conjugate(&block)) * Complex::new(0.5, 0.0);
ham.index_axis_mut(Axis(0), i_r).assign(&herm);
visited[i_r] = true;
} else {
let block_i = ham.index_axis(Axis(0), i_r).to_owned();
let block_j = ham.index_axis(Axis(0), j_r).to_owned();
let avg = (&block_i + &hermitian_conjugate(&block_j)) * Complex::new(0.5, 0.0);
let avg_dag = hermitian_conjugate(&avg);
ham.index_axis_mut(Axis(0), i_r).assign(&avg);
ham.index_axis_mut(Axis(0), j_r).assign(&avg_dag);
visited[i_r] = true;
visited[j_r] = true;
}
}
Ok(())
}
fn hermitian_conjugate(a: &Array2<Complex<f64>>) -> Array2<Complex<f64>> {
a.t().mapv(|x| x.conj())
}
fn normalize3(v: &Array1<f64>) -> Result<Array1<f64>> {
let norm = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
if norm < 1e-14 {
return Err(TbError::Other(
"Cannot normalize a zero-length 3D vector".to_string(),
));
}
Ok(v.mapv(|x| x / norm))
}
fn cross3(a: &Array1<f64>, b: &Array1<f64>) -> Array1<f64> {
arr1(&[
a[1] * b[2] - a[2] * b[1],
a[2] * b[0] - a[0] * b[2],
a[0] * b[1] - a[1] * b[0],
])
}
#[cfg(test)]
mod tests {
use super::*;
use crate::SpinDirection;
use crate::atom_struct::{Atom, AtomType, OrbProj, OrbitalId};
use crate::model::NoRMatrix;
use crate::solve_ham::Solve;
use ndarray::{arr1, array};
#[test]
fn bessel_j_matches_tabulated_values() {
let table: [(isize, f64, f64); 12] = [
(0, 0.0, 1.0),
(1, 0.0, 0.0),
(5, 0.0, 0.0),
(0, 1.0, 0.76519768655796655145),
(1, 1.0, 0.44005058574493351596),
(5, 1.0, 0.00024975773021123443176),
(0, 2.0, 0.22389077914123566805),
(1, 2.0, 0.57672480775687338720),
(2, 2.0, 0.35283402861563771915),
(3, 2.0, 0.12894324947440205110),
(0, 5.0, -0.17759677131433830435),
(10, 5.0, 0.00146780264731047436),
];
for (m, r, expected) in table {
let got = bessel_j(m, r);
assert!(
(got - expected).abs() < 1e-14,
"J_{m}({r}) = {got}, expected {expected}"
);
}
}
#[test]
fn bessel_j_matches_independent_miller_reference() {
let miller = |r: f64, mmax: usize| -> Vec<f64> {
let start = mmax + 20;
let mut next = 0.0_f64; let mut current = 1.0_f64; let mut values = vec![0.0_f64; start + 1];
for k in (0..start).rev() {
values[k] = current;
let k_prev = k as f64;
let prev = if k == 0 {
0.0
} else {
(2.0 * k_prev / r) * current - next
};
next = current;
current = prev;
}
let j0_unscaled = values[0];
let even_sum: f64 = values.iter().step_by(2).skip(1).sum::<f64>();
let scale = 1.0 / (j0_unscaled + 2.0 * even_sum);
values.iter().map(|v| v * scale).collect()
};
for r in [0.3, 1.0, 3.0, 5.0, 8.0] {
let reference = miller(r, 20);
for m in 0..=16 {
let got = bessel_j(m as isize, r);
assert!(
(got - reference[m]).abs() < 1e-11,
"J_{m}({r}) = {got}, Miller reference {}",
reference[m]
);
}
}
let near_zero = bessel_j(1, 7.015586669815619);
assert!(
near_zero.abs() < 1e-12,
"J_1 near its zero must be tiny in absolute value, got {near_zero}"
);
assert!((bessel_j(0, 1e-12) - 1.0).abs() < 3e-16);
assert!((bessel_j(1, 1e-12) - 5e-13).abs() < 1e-25);
}
#[test]
fn bessel_j_satisfies_recurrence_and_negative_order_symmetry() {
for r in [0.3, 0.7, 1.3, 2.5, 4.0, 7.0] {
for m in 1..12 {
let left = bessel_j(m - 1, r) + bessel_j(m + 1, r);
let right = (2.0 * m as f64 / r) * bessel_j(m, r);
assert!(
(left - right).abs() < 1e-12,
"recurrence failed at m={m}, r={r}: {left} vs {right}"
);
}
}
for m in 1..8 {
let expected = if m % 2 == 0 {
bessel_j(m, 1.7)
} else {
-bessel_j(m, 1.7)
};
assert!((bessel_j(-m, 1.7) - expected).abs() < 1e-15);
}
}
#[test]
fn bessel_coeffs_match_single_mode_closed_form() {
let d = array![1.0];
for (amplitude, name) in [
(array![Complex::new(0.4, 0.0)], "linear"),
(array![Complex::new(0.0, 0.4)], "circular"),
(array![Complex::new(0.3, 0.2)], "elliptical"),
] {
let drive = FloquetDrive::with_modes(0.8, vec![LightMode::new(1, amplitude.clone())]);
let r = amplitude[0].norm();
let delta = amplitude[0].arg();
let coeffs = bessel_peierls_coeffs(&d, &drive, -6, 6, 6).unwrap();
for q in -6..=6 {
let expected = Complex::from_polar(1.0, -(q as f64) * std::f64::consts::FRAC_PI_2)
* bessel_j(q, r)
* Complex::from_polar(1.0, (q as f64) * delta);
let got = coeffs[(q + 6) as usize];
assert!(
(got - expected).norm() < 1e-13,
"{name}: C_{q} = {got}, closed form {expected}"
);
}
}
}
#[test]
fn bessel_coeffs_match_time_grid_dft() {
let d = array![1.3, -0.7];
let cases: Vec<FloquetDrive> = vec![
FloquetDrive::with_modes(
1.0,
vec![
LightMode::new(1, array![Complex::new(0.25, 0.0), Complex::new(0.0, 0.25)]),
LightMode::new(2, array![Complex::new(0.1, 0.0), Complex::new(0.05, -0.05)]),
],
),
FloquetDrive::with_modes(
0.7,
vec![
LightMode::new(1, array![Complex::new(0.3, 0.1), Complex::new(-0.1, 0.2)]),
LightMode::new(
-3,
array![Complex::new(0.08, -0.04), Complex::new(0.02, 0.06)],
),
],
),
FloquetDrive::with_modes(
0.5,
vec![
LightMode::new(1, array![Complex::new(0.2, 0.0), Complex::new(0.0, 0.2)]),
LightMode::new(2, array![Complex::new(0.05, 0.0), Complex::new(0.0, -0.05)]),
LightMode::new(
3,
array![Complex::new(0.02, 0.01), Complex::new(0.01, -0.02)],
),
],
),
];
for (case, drive) in cases.iter().enumerate() {
let q_min = -5_isize;
let q_max = 5_isize;
let bessel = bessel_peierls_coeffs(&d, drive, q_min, q_max, 6).unwrap();
let time_grid =
FloquetTimeGrid::new(drive, &FloquetTruncation::new(3, 512), q_min, q_max, 2);
let dft = peierls_fourier_coeffs(&d, q_min, q_max, drive, &time_grid);
for (q, (got, expected)) in bessel.iter().zip(dft.iter()).enumerate() {
assert!(
(got - expected).norm() < 1e-10,
"case {case}, q={}: Bessel {got} vs DFT {expected}",
q_min + q as isize
);
}
}
}
#[test]
fn bessel_coeffs_handle_empty_drive() {
let d = array![0.5];
let drive = FloquetDrive::new(1.0);
let coeffs = bessel_peierls_coeffs(&d, &drive, -3, 3, 6).unwrap();
for q in -3..=3 {
let expected = if q == 0 {
Complex::new(1.0, 0.0)
} else {
Complex::new(0.0, 0.0)
};
assert!((coeffs[(q + 3) as usize] - expected).norm() < 1e-15);
}
}
#[test]
fn harmonic_cache_bessel_matches_time_grid_and_dedupes_links() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [0.3, 0.0]];
let mut model = Model::<true, 2>::tb_model(lat, orb, None).unwrap();
model.add_hop(-1.0, 0, 0, &array![1, 0], None);
model.add_hop(-0.5, 0, 1, &array![0, 1], None);
model.add_hop(
Complex::new(0.1, 0.2),
0,
1,
&array![1, 1],
SpinDirection::X,
);
let mut distinct_d = Vec::<Vec<u64>>::new();
for i_r in 0..model.hamR.nrows() {
let r_vec = model.hamR.row(i_r);
for i in 0..model.nsta() {
for j in 0..model.nsta() {
if model.ham[[i_r, i, j]].norm_sqr() == 0.0 {
continue;
}
let d_cart = model.link_displacement_cartesian(
i % model.norb(),
j % model.norb(),
&r_vec,
);
let key: Vec<u64> = d_cart.iter().map(|value| value.to_bits()).collect();
if !distinct_d.contains(&key) {
distinct_d.push(key);
}
}
}
}
assert_eq!(
distinct_d.len(),
6,
"dedup premise: expected 6 distinct link displacements, found {}",
distinct_d.len()
);
let drive = FloquetDrive::with_modes(
0.8,
vec![
LightMode::new(1, array![Complex::new(0.2, 0.0), Complex::new(0.0, 0.2)]),
LightMode::new(2, array![Complex::new(0.05, -0.05), Complex::new(0.0, 0.0)]),
],
);
let trunc = FloquetTruncation::new(2, 512);
let time_grid_cache =
model.floquet_harmonic_cache(&drive, &trunc, -4, 4, &PeierlsFourierMethod::TimeGrid);
let bessel_cache = model.floquet_harmonic_cache(
&drive,
&trunc,
-4,
4,
&PeierlsFourierMethod::Bessel { cutoff_margin: 6 },
);
assert_eq!(time_grid_cache.blocks.dim(), bessel_cache.blocks.dim());
for (a, b) in time_grid_cache
.blocks
.iter()
.zip(bessel_cache.blocks.iter())
{
assert!(
(a - b).norm() < 1e-10,
"Bessel cache {b} vs time-grid cache {a}"
);
}
}
#[test]
fn harmonic_cache_bessel_falls_back_for_large_amplitudes() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [10.0, 0.0]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.add_hop(-1.0, 0, 0, &array![1, 0], None);
model.add_hop(-0.5, 0, 1, &array![0, 1], None);
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
1,
array![Complex::new(0.9, 0.0), Complex::new(0.0, 0.0)],
)],
);
let trunc = FloquetTruncation::new(1, 512);
let time_grid_cache =
model.floquet_harmonic_cache(&drive, &trunc, -3, 3, &PeierlsFourierMethod::TimeGrid);
let bessel_cache = model.floquet_harmonic_cache(
&drive,
&trunc,
-3,
3,
&PeierlsFourierMethod::Bessel { cutoff_margin: 6 },
);
for (a, b) in time_grid_cache
.blocks
.iter()
.zip(bessel_cache.blocks.iter())
{
assert!(
(a - b).norm() < 1e-10,
"fallback Bessel cache {b} vs time-grid cache {a}"
);
}
}
#[test]
fn harmonic_cache_bessel_fallback_uses_alias_free_grid() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [55.555_555_555_555_56, 0.0]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.add_hop(-1.0, 0, 1, &array![0, 1], None);
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
100,
array![Complex::new(0.9, 0.0), Complex::new(0.0, 0.0)],
)],
);
let oracle = model.floquet_harmonic_cache(
&drive,
&FloquetTruncation::new(100, 65536),
-10,
10,
&PeierlsFourierMethod::TimeGrid,
);
let bessel_cache = model.floquet_harmonic_cache(
&drive,
&FloquetTruncation::new(100, 512),
-10,
10,
&PeierlsFourierMethod::Bessel { cutoff_margin: 6 },
);
for (a, b) in oracle.blocks.iter().zip(bessel_cache.blocks.iter()) {
assert!(
(a - b).norm() < 1e-10,
"65536-point oracle {a} vs adaptive fallback {b}"
);
}
let i_r_link = find_R(&model.hamR, &array![0, 1]).unwrap();
let c0 = bessel_cache.blocks[[bessel_cache.q_index(0), i_r_link, 0, 1]];
assert!(
(c0 - Complex::new(-bessel_j(0, 50.0), 0.0)).norm() < 1e-10,
"C_0 on the r = 50 link should be -J_0(50) (t = -1), got {c0}"
);
}
fn commutator_test_blocks<const DIM: usize>(
model: &Model<false, DIM, NoRMatrix>,
drive: &FloquetDrive,
) -> (
FloquetHarmonicCache,
Vec<Array2<Complex<f64>>>,
Vec<Array2<Complex<f64>>>,
) {
let trunc = FloquetTruncation::new(1, 512);
let cache = model.floquet_harmonic_cache(
drive,
&trunc,
-1,
1,
&PeierlsFourierMethod::Bessel { cutoff_margin: 6 },
);
let n_r = model.hamR.nrows();
let q1 = cache.q_index(1);
let qm1 = cache.q_index(-1);
let a_blocks = (0..n_r)
.map(|i_r| cache.blocks.slice(s![q1, i_r, .., ..]).to_owned())
.collect();
let b_blocks = (0..n_r)
.map(|i_r| cache.blocks.slice(s![qm1, i_r, .., ..]).to_owned())
.collect();
(cache, a_blocks, b_blocks)
}
#[test]
fn real_space_commutator_matches_k_space_commutator() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [0.35, 0.2]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1, 0], None);
model.set_hop(-0.3, 0, 1, &array![0, 1], None);
model.set_hop(Complex::new(0.1, -0.2), 1, 1, &array![1, 1], None);
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
1,
array![Complex::new(0.3, 0.0), Complex::new(0.0, 0.3)],
)],
);
let (cache, a_blocks, b_blocks) = commutator_test_blocks(&model, &drive);
let (comm_blocks, comm_r) =
real_space_commutator(&a_blocks, &b_blocks, &model.hamR).unwrap();
let nsta = model.nsta();
let mut oracle_scale = 0.0_f64;
for k in [[0.0, 0.0], [0.123, 0.321], [0.5, 0.5], [0.877, 0.111]] {
let kvec = array![k[0], k[1]];
let a_k = model.floquet_cached_harmonic_onek(&kvec, 1, Gauge::Lattice, &cache);
let b_k = model.floquet_cached_harmonic_onek(&kvec, -1, Gauge::Lattice, &cache);
let oracle = a_k.dot(&b_k) - b_k.dot(&a_k);
oracle_scale = oracle_scale.max(oracle.iter().fold(0.0, |m, c| m.max(c.norm())));
let mut from_rs = Array2::<Complex<f64>>::zeros((nsta, nsta));
for (i_r, row) in comm_r.outer_iter().enumerate() {
let mut phase_arg = 0.0;
for a in 0..2 {
phase_arg += row[a] as f64 * kvec[a];
}
let phase = Complex::new(0.0, TAU * phase_arg).exp();
from_rs.scaled_add(phase, &comm_blocks[i_r]);
}
for (a, b) in from_rs.iter().zip(oracle.iter()) {
assert!(
(a - b).norm() < 1e-12,
"k = {k:?}: real-space {a} vs k-space {b}"
);
}
}
assert!(
oracle_scale > 1e-6,
"oracle must be non-trivial (circular 2D drive); otherwise \
the comparison is vacuous, got {oracle_scale}"
);
}
#[test]
fn real_space_commutator_vanishes_for_linear_polarization() {
let lat = array![[1.0]];
let orb = array![[0.0], [0.35]];
let mut model = Model::<false, 1>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1], None);
model.set_hop(Complex::new(-0.3, 0.1), 0, 1, &array![1], None);
model.set_hop(Complex::new(0.1, -0.2), 1, 1, &array![2], None);
let drive =
FloquetDrive::with_modes(1.0, vec![LightMode::new(1, array![Complex::new(0.4, 0.2)])]);
let (_cache, a_blocks, b_blocks) = commutator_test_blocks(&model, &drive);
let (comm_blocks, _comm_r) =
real_space_commutator(&a_blocks, &b_blocks, &model.hamR).unwrap();
for block in &comm_blocks {
let max = block.iter().fold(0.0_f64, |m, c| m.max(c.norm()));
assert!(
max < 1e-15,
"linear-polarization commutator must vanish exactly, got {max}"
);
}
}
#[test]
fn real_space_commutator_support_and_hermiticity() {
let lat = array![[1.0]];
let orb = array![[0.0], [0.35]];
let mut model = Model::<false, 1>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1], None);
model.set_hop(-0.3, 0, 1, &array![1], None);
model.set_hop(Complex::new(0.1, -0.2), 1, 1, &array![2], None);
let drive =
FloquetDrive::with_modes(1.0, vec![LightMode::new(1, array![Complex::new(0.4, 0.2)])]);
let (_cache, a_blocks, b_blocks) = commutator_test_blocks(&model, &drive);
let (comm_blocks, comm_r) =
real_space_commutator(&a_blocks, &b_blocks, &model.hamR).unwrap();
let expected: Vec<Vec<isize>> = (-4..=4).map(|r| vec![r]).collect();
let got: Vec<Vec<isize>> = comm_r.outer_iter().map(|row| row.to_vec()).collect();
assert_eq!(
got, expected,
"commutator support must be the Minkowski sum"
);
for i in 0..comm_r.nrows() {
let j = find_R(&comm_r, &comm_r.row(i).mapv(|v| -v)).unwrap();
let conj = hermitian_conjugate(&comm_blocks[j]);
for (a, b) in comm_blocks[i].iter().zip(conj.iter()) {
assert!(
(a - b).norm() < 1e-15,
"comm(R) != comm(−R)† at R = {:?}",
comm_r.row(i).to_vec()
);
}
}
}
#[test]
fn real_space_commutator_scalar_blocks_vanish() {
let lat = array![[1.0]];
let orb = array![[0.0]];
let mut model = Model::<false, 1>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1], None);
let drive =
FloquetDrive::with_modes(1.0, vec![LightMode::new(1, array![Complex::new(0.4, 0.2)])]);
let (_cache, a_blocks, b_blocks) = commutator_test_blocks(&model, &drive);
let (comm_blocks, _comm_r) =
real_space_commutator(&a_blocks, &b_blocks, &model.hamR).unwrap();
for block in &comm_blocks {
assert!(
block[[0, 0]].norm() < 1e-15,
"scalar commutator must vanish, got {}",
block[[0, 0]]
);
}
}
#[test]
fn floquet_effective_model_bessel_matches_legacy_bands() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [0.35, 0.2]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1, 0], None);
model.set_hop(-0.3, 0, 1, &array![0, 1], None);
model.set_hop(Complex::new(0.1, -0.2), 1, 1, &array![1, 1], None);
let drive = FloquetDrive::with_modes(
1.0,
vec![
LightMode::new(1, array![Complex::new(0.3, 0.0), Complex::new(0.0, 0.3)]),
LightMode::new(2, array![Complex::new(0.05, -0.05), Complex::new(0.0, 0.0)]),
],
);
let trunc = FloquetTruncation::new(2, 4096);
let options = FloquetEffectiveOptions::new().with_q_max(4);
let bessel = model
.floquet_effective_model(&drive, &trunc, Some(&options))
.unwrap();
let legacy = model
.floquet_effective_model_legacy(
&drive,
&trunc,
[64, 64],
Some(&options.with_target_hamR(bessel.hamR.clone())),
)
.unwrap();
for k in [[0.1, 0.2], [0.5, 0.5], [0.9, 0.7]] {
let kvec = array![k[0], k[1]];
let e_b = eigvalsh_v(&bessel.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
let e_l = eigvalsh_v(&legacy.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
for (a, b) in e_b.iter().zip(e_l.iter()) {
assert!((a - b).abs() < 1e-8, "k = {k:?}: Bessel {a} vs legacy {b}");
}
}
}
#[test]
fn floquet_effective_model_bessel_order_zero() {
let lat = array![[1.0]];
let orb = array![[0.0], [0.35]];
let mut model = Model::<false, 1>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1], None);
model.set_hop(Complex::new(-0.3, 0.1), 0, 1, &array![1], None);
let drive =
FloquetDrive::with_modes(1.0, vec![LightMode::new(1, array![Complex::new(0.4, 0.2)])]);
let trunc = FloquetTruncation::new(1, 4096);
let bessel = model
.floquet_effective_model(
&drive,
&trunc,
Some(&FloquetEffectiveOptions::new().with_order(0)),
)
.unwrap();
assert_eq!(bessel.hamR.nrows(), model.hamR.nrows());
let legacy = model
.floquet_effective_model_legacy(
&drive,
&trunc,
[64],
Some(&FloquetEffectiveOptions::new().with_order(0)),
)
.unwrap();
for k in [0.0, 0.23, 0.5, 0.71] {
let kvec = array![k];
let e_b = eigvalsh_v(&bessel.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
let e_l = eigvalsh_v(&legacy.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
for (a, b) in e_b.iter().zip(e_l.iter()) {
assert!((a - b).abs() < 1e-8, "k = {k}: Bessel {a} vs legacy {b}");
}
}
}
#[test]
fn floquet_effective_model_bessel_support_and_hermiticity() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [0.35, 0.2]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1, 0], None);
model.set_hop(-0.3, 0, 1, &array![0, 1], None);
model.set_hop(Complex::new(0.1, -0.2), 1, 1, &array![1, 1], None);
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
1,
array![Complex::new(0.3, 0.0), Complex::new(0.0, 0.3)],
)],
);
let bessel = model
.floquet_effective_model(&drive, &FloquetTruncation::new(1, 512), None)
.unwrap();
let mut expected = std::collections::BTreeSet::<Vec<isize>>::new();
for r1 in model.hamR.outer_iter() {
for r2 in model.hamR.outer_iter() {
expected.insert(r1.iter().zip(r2.iter()).map(|(a, b)| a + b).collect());
}
}
let got: Vec<Vec<isize>> = bessel.hamR.outer_iter().map(|row| row.to_vec()).collect();
assert_eq!(got, Vec::from_iter(expected), "support mismatch");
for i in 0..bessel.hamR.nrows() {
let j = find_R(&bessel.hamR, &bessel.hamR.row(i).mapv(|v| -v)).unwrap();
let conj = hermitian_conjugate(&bessel.ham.index_axis(Axis(0), j).to_owned());
for (a, b) in bessel.ham.index_axis(Axis(0), i).iter().zip(conj.iter()) {
assert!((a - b).norm() < 1e-15, "T(R) != T(−R)† at row {i}");
}
}
}
#[test]
fn floquet_effective_model_bessel_rejects_invalid_options() {
let lat = array![[1.0]];
let orb = array![[0.0]];
let mut model = Model::<false, 1>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1], None);
let drive =
FloquetDrive::with_modes(1.0, vec![LightMode::new(1, array![Complex::new(0.4, 0.2)])]);
let trunc = FloquetTruncation::new(1, 512);
assert!(
model
.floquet_effective_model(
&drive,
&trunc,
Some(&FloquetEffectiveOptions::new().with_order(2))
)
.is_err()
);
let target = array![[-1], [0], [1]];
assert!(
model
.floquet_effective_model(
&drive,
&trunc,
Some(&FloquetEffectiveOptions::new().with_target_hamR(target))
)
.is_err()
);
assert!(
model
.floquet_effective_model(
&drive,
&trunc,
Some(&FloquetEffectiveOptions::new().with_q_max(-1))
)
.is_err()
);
}
#[test]
fn floquet_effective_model_spinful_matches_legacy() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [0.35, 0.2]];
let mut model = Model::<true, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1, 0], None);
model.set_hop(-0.3, 0, 1, &array![0, 1], SpinDirection::X);
model.set_hop(
Complex::new(0.1, -0.2),
1,
1,
&array![1, 1],
SpinDirection::Z,
);
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
1,
array![Complex::new(0.3, 0.0), Complex::new(0.0, 0.3)],
)],
);
let trunc = FloquetTruncation::new(1, 4096);
let bessel = model.floquet_effective_model(&drive, &trunc, None).unwrap();
let legacy = model
.floquet_effective_model_legacy(
&drive,
&trunc,
[32, 32],
Some(&FloquetEffectiveOptions::new().with_target_hamR(bessel.hamR.clone())),
)
.unwrap();
for k in [[0.1, 0.2], [0.5, 0.5]] {
let kvec = array![k[0], k[1]];
let e_b = eigvalsh_v(&bessel.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
let e_l = eigvalsh_v(&legacy.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
for (a, b) in e_b.iter().zip(e_l.iter()) {
assert!((a - b).abs() < 1e-8, "k = {k:?}: Bessel {a} vs legacy {b}");
}
}
}
#[test]
fn floquet_effective_model_no_drive_matches_static_bands() {
let model = chain_model();
let drive = FloquetDrive::new(0.9);
let trunc = FloquetTruncation::new(2, 64);
let effective = model.floquet_effective_model(&drive, &trunc, None).unwrap();
let k = arr1(&[0.271]);
for gauge in [Gauge::Lattice, Gauge::Atom] {
let from_effective = effective.gen_ham(&k, gauge);
let from_static = model.gen_ham(&k, gauge);
let mut max_diff = 0.0f64;
for i in 0..from_effective.nrows() {
for j in 0..from_effective.ncols() {
max_diff = max_diff.max((from_effective[[i, j]] - from_static[[i, j]]).norm());
}
}
assert!(
max_diff < 1e-12,
"no-drive effective mismatch in {gauge:?}: {max_diff:e}"
);
}
let got: Vec<Vec<isize>> = effective
.hamR
.outer_iter()
.map(|row| row.to_vec())
.collect();
let expected: Vec<Vec<isize>> = (-2..=2).map(|r| vec![r]).collect();
assert_eq!(
got, expected,
"empty-drive support must be the Minkowski union"
);
}
#[test]
fn floquet_effective_model_bessel_perf_smoke() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [0.35, 0.2]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0, 0, 0, &array![1, 0], None);
model.set_hop(-0.3, 0, 1, &array![0, 1], None);
model.set_hop(Complex::new(0.1, -0.2), 1, 1, &array![1, 1], None);
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
1,
array![Complex::new(0.3, 0.0), Complex::new(0.0, 0.3)],
)],
);
let trunc = FloquetTruncation::new(1, 512);
let warmup = model.floquet_effective_model(&drive, &trunc, None).unwrap();
let legacy_options = FloquetEffectiveOptions::new().with_target_hamR(warmup.hamR.clone());
let _ = model
.floquet_effective_model_legacy(&drive, &trunc, [128, 128], Some(&legacy_options))
.unwrap();
let start = std::time::Instant::now();
let bessel = model.floquet_effective_model(&drive, &trunc, None).unwrap();
let mut t_bessel = start.elapsed();
for _ in 0..2 {
let start = std::time::Instant::now();
let _ = model.floquet_effective_model(&drive, &trunc, None).unwrap();
t_bessel = t_bessel.min(start.elapsed());
}
let start = std::time::Instant::now();
let legacy = model
.floquet_effective_model_legacy(&drive, &trunc, [128, 128], Some(&legacy_options))
.unwrap();
let mut t_legacy = start.elapsed();
for _ in 0..2 {
let start = std::time::Instant::now();
let _ = model
.floquet_effective_model_legacy(&drive, &trunc, [128, 128], Some(&legacy_options))
.unwrap();
t_legacy = t_legacy.min(start.elapsed());
}
let kvec = array![0.37, 0.19];
let e_b = eigvalsh_v(&bessel.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
let e_l = eigvalsh_v(&legacy.gen_ham(&kvec, Gauge::Lattice), UPLO::Lower);
for (a, b) in e_b.iter().zip(e_l.iter()) {
assert!((a - b).abs() < 1e-8, "Bessel {a} vs legacy {b}");
}
let ratio = t_legacy.as_secs_f64() / t_bessel.as_secs_f64().max(1e-9);
eprintln!("perf smoke: bessel {t_bessel:?} vs legacy {t_legacy:?} ({ratio:.0}x)");
let speedup = t_legacy.as_secs_f64()
/ t_bessel
.max(std::time::Duration::from_micros(100))
.as_secs_f64();
assert!(
speedup > 5.0,
"real-space path should be far faster than the k-space path ({speedup:.1}x)"
);
}
fn graphene_model(j: f64) -> Model<false, 2, NoRMatrix> {
let lat = array![[3f64.sqrt() / 2.0, 0.5], [3f64.sqrt() / 2.0, -0.5],];
let orb = array![[0.0, 0.0], [0.0, 0.0]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
for r in [[1, -1], [-1, 0], [0, 1]] {
model.set_hop(j, 0, 1, &array![r[0], r[1]], None);
}
model
}
fn graphene_q_n_lit(n: isize, k: &[f64; 2], j: f64, alpha: f64) -> Complex<f64> {
let e_int = [[1, -1], [-1, 0], [0, 1]];
let mut sum = Complex::new(0.0, 0.0);
for (l, r) in e_int.iter().enumerate() {
let phase = -TAU * (k[0] * r[0] as f64 + k[1] * r[1] as f64)
+ TAU * (n as f64) * (l as f64) / 3.0;
sum += Complex::new(0.0, phase).exp();
}
j * bessel_j(n, alpha) * sum
}
fn graphene_circular_drive(alpha: f64) -> FloquetDrive {
FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
1,
array![Complex::new(alpha, 0.0), Complex::new(0.0, alpha)],
)],
)
}
#[test]
fn graphene_harmonics_match_literature_fourier_components() {
let j = -1.0;
let model = graphene_model(j);
let trunc = FloquetTruncation::new(4, 512);
for alpha in [0.3, 0.8] {
let drive = graphene_circular_drive(alpha);
let cache = model.floquet_harmonic_cache(
&drive,
&trunc,
-4,
4,
&PeierlsFourierMethod::Bessel { cutoff_margin: 6 },
);
for k in [[0.13, 0.21], [0.5, 0.5], [0.87, 0.11]] {
let kvec = array![k[0], k[1]];
let k_neg = [-k[0], -k[1]];
for q in -4..=4 {
let hq = model.floquet_cached_harmonic_onek(&kvec, q, Gauge::Lattice, &cache);
let lit = array![
[
Complex::new(0.0, 0.0),
graphene_q_n_lit(q, &k_neg, j, alpha),
],
[
graphene_q_n_lit(-q, &k_neg, j, alpha).conj(),
Complex::new(0.0, 0.0),
],
];
for (a, b) in hq.iter().zip(lit.iter()) {
assert!(
(a - b).norm() < 1e-12,
"alpha = {alpha}, k = {k:?}, q = {q}: {a} vs {b}"
);
}
}
}
}
}
#[test]
fn graphene_dirac_gap_matches_exact_rotating_frame() {
let j = -1.0;
let model = graphene_model(j);
let alpha = 0.2;
let drive = FloquetDrive::with_modes(
5.0,
vec![LightMode::new(
1,
array![Complex::new(alpha, 0.0), Complex::new(0.0, alpha)],
)],
);
let w = drive.omega0_ev;
let g = 1.5 * j.abs() * alpha;
let delta_exact = (w * w + 4.0 * g * g).sqrt() - w;
let k = array![1.0 / 3.0, 1.0 / 3.0];
let n_cut = 8;
let k_eff = -(2.0 * j * j / w)
* (1..=n_cut)
.map(|n| bessel_j(n, alpha).powi(2) / (n as f64) * (TAU * (n as f64) / 3.0).sin())
.sum::<f64>();
let d_z_k = -3.0 * 3f64.sqrt() * k_eff;
let mut outer_previous = f64::INFINITY;
for n_max in [4, 8, 12] {
let trunc = FloquetTruncation::new(n_max, 512);
let hf = model
.floquet_ham_onek(&k, &drive, &trunc, Gauge::Lattice)
.unwrap();
let e = eigvalsh_v(&hf, UPLO::Lower);
let outer = e
.iter()
.map(|x| ((*x + w / 2.0).rem_euclid(w) - w / 2.0).abs())
.fold(0.0_f64, f64::max);
let tol = if n_max <= 4 { 2e-3 } else { 1e-3 };
assert!(
(outer - delta_exact / 2.0).abs() < tol,
"n_max = {n_max}: outer branch {outer} vs exact {:.6}",
delta_exact / 2.0
);
assert!(
(outer - d_z_k.abs()).abs() < 5e-3,
"n_max = {n_max}: outer branch {outer} vs van Vleck mass {d_z_k}"
);
assert!(
outer <= outer_previous + 1e-12,
"outer branch must converge downward: {outer} vs previous {outer_previous}"
);
outer_previous = outer;
}
}
#[test]
fn graphene_haldane_mass_matches_full_bessel_series() {
let j = -1.0;
let model = graphene_model(j);
let alpha = 0.5;
let drive = FloquetDrive::with_modes(
5.0,
vec![LightMode::new(
1,
array![Complex::new(alpha, 0.0), Complex::new(0.0, alpha)],
)],
);
let w = drive.omega0_ev;
let trunc = FloquetTruncation::new(4, 512);
let b_int = [[1, 1], [1, -2], [-2, 1]];
let k_eff_series = |n_cut: isize| -> f64 {
-(2.0 * j * j / w)
* (1..=n_cut)
.map(|n| {
bessel_j(n, alpha).powi(2) / (n as f64) * (TAU * (n as f64) / 3.0).sin()
})
.sum::<f64>()
};
let d_z_lit = |k: &[f64; 2], n_cut: isize| -> f64 {
2.0 * k_eff_series(n_cut)
* b_int
.iter()
.map(|b| (TAU * (k[0] * b[0] as f64 + k[1] * b[1] as f64)).sin())
.sum::<f64>()
};
let eff = model
.floquet_effective_model(
&drive,
&trunc,
Some(&FloquetEffectiveOptions::new().with_q_max(8)),
)
.unwrap();
let eff2 = model
.floquet_effective_model(
&drive,
&trunc,
Some(&FloquetEffectiveOptions::new().with_q_max(2)),
)
.unwrap();
for k in [[0.13, 0.07], [0.23, 0.11], [-0.07, 0.31]] {
let kvec = array![k[0], k[1]];
let h = eff.gen_ham(&kvec, Gauge::Lattice);
let h2 = eff2.gen_ham(&kvec, Gauge::Lattice);
let d_z = ((h[[0, 0]] - h[[1, 1]]).re) / 2.0;
let d_z_2 = ((h2[[0, 0]] - h2[[1, 1]]).re) / 2.0;
assert!(
(d_z - d_z_lit(&k, 8)).abs() < 1e-8,
"k = {k:?}: d_z {d_z} vs series {expect}",
expect = d_z_lit(&k, 8)
);
assert!(
(d_z_2 - d_z_lit(&k, 2)).abs() < 1e-8,
"k = {k:?}: q_max = 2: d_z {d_z_2} vs series {}",
d_z_lit(&k, 2)
);
assert!(
(d_z - d_z_2).abs() < 1e-5,
"k = {k:?}: truncation tail {d_z} vs {d_z_2}"
);
}
}
#[test]
fn graphene_order_zero_matches_renormalized_nn_hopping() {
let j = -1.0;
let model = graphene_model(j);
let trunc = FloquetTruncation::new(1, 512);
let alpha = 0.6;
let drive = graphene_circular_drive(alpha);
let eff = model
.floquet_effective_model(
&drive,
&trunc,
Some(&FloquetEffectiveOptions::new().with_order(0)),
)
.unwrap();
for k in [[0.13, 0.21], [0.5, 0.5], [0.87, 0.11]] {
let kvec = array![k[0], k[1]];
let h0 = eff.gen_ham(&kvec, Gauge::Lattice);
let lit = array![
[
Complex::new(0.0, 0.0),
graphene_q_n_lit(0, &[-k[0], -k[1]], j, alpha),
],
[
graphene_q_n_lit(0, &[-k[0], -k[1]], j, alpha).conj(),
Complex::new(0.0, 0.0),
],
];
for (a, b) in h0.iter().zip(lit.iter()) {
assert!(
(a - b).norm() < 1e-12,
"alpha = {alpha}, k = {k:?}: {a} vs {b}"
);
}
}
}
#[test]
fn bessel_coeffs_reject_large_amplitudes_and_bad_ranges() {
let d = array![60.0];
let drive =
FloquetDrive::with_modes(1.0, vec![LightMode::new(1, array![Complex::new(1.0, 0.0)])]);
assert!(bessel_peierls_coeffs(&d, &drive, -4, 4, 6).is_err());
let zero_l =
FloquetDrive::with_modes(1.0, vec![LightMode::new(0, array![Complex::new(0.1, 0.0)])]);
let out_of_range = bessel_peierls_coeffs(&d, &zero_l, 5, 7, 6).unwrap();
for q in 5..=7 {
assert!((out_of_range[(q - 5) as usize]).norm() == 0.0);
}
assert!(bessel_peierls_coeffs(&d, &drive, 3, 2, 6).is_err());
assert!(bessel_peierls_coeffs(&d, &drive, -4, 4, 49).is_err());
assert!(bessel_peierls_coeffs(&d, &drive, -4, 4, -1).is_err());
}
#[test]
fn bessel_coeffs_large_harmonics_stay_within_error_budget() {
let d = array![1.0, 0.0];
let drive = FloquetDrive::with_modes(
1.0,
vec![
LightMode::new(400, array![Complex::new(8.0, 0.0), Complex::new(0.0, 0.0)]),
LightMode::new(-400, array![Complex::new(8.0, 0.0), Complex::new(0.0, 0.0)]),
],
);
let coeffs = bessel_peierls_coeffs(&d, &drive, -10, 10, 0).unwrap();
for q in -10..=10 {
let expected = if q == 0 { bessel_j(0, 16.0) } else { 0.0 };
let got = coeffs[(q + 10) as usize];
assert!(
(got - Complex::new(expected, 0.0)).norm() < 1e-12,
"q = {q}: got {got}, expected {expected}"
);
}
}
#[test]
fn bessel_coeffs_window_edge_overflows_are_skipped() {
let d = array![1.0, 0.0];
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(
400,
array![Complex::new(8.0, 0.0), Complex::new(0.0, 0.0)],
)],
);
let q = isize::MAX - 23_999;
let coeffs = bessel_peierls_coeffs(&d, &drive, q, q, 48).unwrap();
assert!(coeffs[0].norm() == 0.0);
}
fn chain_model() -> Model<false, 1, NoRMatrix> {
let lat = array![[1.0]];
let orb = array![[0.0]];
let mut model = Model::<false, 1>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0_f64, 0, 0, &arr1(&[1isize]), None);
model
}
fn metadata_model() -> Model<false, 1, NoRMatrix> {
let lat = array![[1.0]];
let orb = array![[0.0], [0.0], [0.35]];
let atoms = vec![
Atom::with_orbitals(
arr1(&[0.0]),
AtomType::C,
[OrbitalId::new(0), OrbitalId::new(1)],
),
Atom::with_orbitals(arr1(&[0.35]), AtomType::O, [OrbitalId::new(2)]),
];
let mut model = Model::<false, 1>::tb_model(lat, orb, Some(atoms)).unwrap();
model.orb_projection = vec![OrbProj::s, OrbProj::px, OrbProj::py];
model.set_hop(-0.8_f64, 0, 2, &arr1(&[1isize]), None);
model
}
fn assert_same_atom_metadata(expected: &Atom, actual: &Atom) {
assert_eq!(actual.position(), expected.position());
assert_eq!(actual.norb(), expected.norb());
assert_eq!(actual.atom_type(), expected.atom_type());
}
#[test]
fn floquet_no_drive_static_replicas() {
let model = chain_model();
let k = arr1(&[0.17]);
let drive = FloquetDrive::new(0.7);
let trunc = FloquetTruncation::new(1, 64);
let bands = model
.floquet_band_onek(&k, &drive, &trunc, Gauge::Atom)
.unwrap();
let e0 = model.solve_band_onek(&k)[0];
let mut expected = vec![e0 - 0.7, e0, e0 + 0.7];
expected.sort_by(|a, b| a.partial_cmp(b).unwrap());
for (a, b) in bands.iter().zip(expected.iter()) {
assert!((a - b).abs() < 1e-12, "got {a}, expected {b}");
}
}
#[test]
fn floquet_hamiltonian_is_hermitian() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0_f64, 0, 0, &arr1(&[1isize, 0]), None);
model.set_hop(-0.7_f64, 0, 0, &arr1(&[0isize, 1]), None);
let drive = FloquetDrive::with_modes(
0.9,
vec![LightMode::new(
1,
arr1(&[Complex::new(0.11, 0.0), Complex::new(0.0, 0.07)]),
)],
);
let trunc = FloquetTruncation::new(2, 512);
let k = arr1(&[0.13, 0.29]);
let hf = model
.floquet_ham_onek(&k, &drive, &trunc, Gauge::Atom)
.unwrap();
let mut max_diff = 0.0f64;
for i in 0..hf.nrows() {
for j in 0..hf.ncols() {
max_diff = max_diff.max((hf[[i, j]] - hf[[j, i]].conj()).norm());
}
}
assert!(max_diff < 1e-11, "max hermiticity error = {max_diff:e}");
}
#[test]
fn floquet_model_matches_onek_construction() {
let lat = array![[1.0, 0.0], [0.2, 1.1]];
let orb = array![[0.0, 0.0], [0.31, 0.17]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(0.2_f64, 0, 1, &arr1(&[0isize, 0]), None);
model.set_hop(-1.0_f64, 0, 0, &arr1(&[1isize, 0]), None);
model.set_hop(-0.6_f64, 1, 1, &arr1(&[0isize, 1]), None);
let drive = FloquetDrive::with_modes(
0.8,
vec![LightMode::new(
1,
arr1(&[Complex::new(0.13, 0.0), Complex::new(0.0, 0.09)]),
)],
);
let trunc = FloquetTruncation::new(1, 256);
let k = arr1(&[0.23, 0.31]);
let floquet_model = model.floquet_model(&drive, &trunc).unwrap();
assert_eq!(floquet_model.nsta(), model.nsta() * trunc.n_sector());
assert_eq!(floquet_model.hamR, model.hamR);
for gauge in [Gauge::Lattice, Gauge::Atom] {
let from_model = floquet_model.gen_ham(&k, gauge);
let from_onek = model.floquet_ham_onek(&k, &drive, &trunc, gauge).unwrap();
let mut max_diff = 0.0f64;
for i in 0..from_model.nrows() {
for j in 0..from_model.ncols() {
max_diff = max_diff.max((from_model[[i, j]] - from_onek[[i, j]]).norm());
}
}
assert!(
max_diff < 1e-12,
"floquet_model mismatch in {gauge:?}: {max_diff:e}"
);
}
}
#[test]
fn floquet_model_preserves_spinful_layout() {
let lat = array![[1.0, 0.0], [0.0, 1.0]];
let orb = array![[0.0, 0.0], [0.27, 0.19]];
let mut model = Model::<true, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(0.3_f64, 0, 0, &arr1(&[0isize, 0]), crate::SpinDirection::Z);
model.add_hop(0.2_f64, 0, 0, &arr1(&[0isize, 0]), crate::SpinDirection::X);
model.set_hop(-0.9_f64, 0, 1, &arr1(&[1isize, 0]), None);
model.set_hop(-0.4_f64, 1, 1, &arr1(&[0isize, 1]), None);
let drive = FloquetDrive::with_modes(
0.6,
vec![LightMode::new(
1,
arr1(&[Complex::new(0.07, 0.0), Complex::new(0.0, 0.05)]),
)],
);
let trunc = FloquetTruncation::new(1, 256);
let k = arr1(&[0.17, 0.29]);
let floquet_model = model.floquet_model(&drive, &trunc).unwrap();
assert_eq!(floquet_model.norb(), model.norb() * trunc.n_sector());
assert_eq!(floquet_model.nsta(), model.nsta() * trunc.n_sector());
for gauge in [Gauge::Lattice, Gauge::Atom] {
let from_model = floquet_model.gen_ham(&k, gauge);
let from_onek = model.floquet_ham_onek(&k, &drive, &trunc, gauge).unwrap();
let mut max_diff = 0.0f64;
for i in 0..from_model.nrows() {
for j in 0..from_model.ncols() {
max_diff = max_diff.max((from_model[[i, j]] - from_onek[[i, j]]).norm());
}
}
assert!(
max_diff < 1e-12,
"spinful floquet_model mismatch in {gauge:?}: {max_diff:e}"
);
}
}
#[test]
fn floquet_models_preserve_atom_metadata() {
let model = metadata_model();
let drive = FloquetDrive::new(1.2);
let trunc = FloquetTruncation::new(1, 32);
let sambe = model.floquet_model(&drive, &trunc).unwrap();
assert_eq!(sambe.natom(), model.natom() * trunc.n_sector());
for sector in 0..trunc.n_sector() {
for i_atom in 0..model.natom() {
assert_same_atom_metadata(
&model.atoms[i_atom],
&sambe.atoms[sector * model.natom() + i_atom],
);
}
}
let expected_projection: Vec<OrbProj> = (0..trunc.n_sector())
.flat_map(|_| model.orb_projection.iter().copied())
.collect();
assert_eq!(sambe.orb_projection, expected_projection);
let effective = model.floquet_effective_model(&drive, &trunc, None).unwrap();
assert_eq!(effective.natom(), model.natom());
for i_atom in 0..model.natom() {
assert_same_atom_metadata(&model.atoms[i_atom], &effective.atoms[i_atom]);
}
assert_eq!(effective.orb_projection, model.orb_projection);
}
#[test]
fn floquet_effective_rejects_non_hermitian_target_range() {
let model = chain_model();
let drive = FloquetDrive::new(1.0);
let trunc = FloquetTruncation::new(1, 32);
let options = FloquetEffectiveOptions::new().with_target_hamR(array![[0isize], [1isize]]);
let err = model
.floquet_effective_model_legacy(&drive, &trunc, [8], Some(&options))
.unwrap_err();
match err {
TbError::MissingHermitianConjugateHopping { r } => {
assert_eq!(r, arr1(&[1isize]));
}
other => panic!("unexpected error: {other}"),
}
}
#[test]
fn floquet_effective_rejects_duplicate_target_vectors() {
let model = chain_model();
let drive = FloquetDrive::new(1.0);
let trunc = FloquetTruncation::new(1, 32);
let options = FloquetEffectiveOptions::new().with_target_hamR(array![[0isize], [0isize]]);
let err = model
.floquet_effective_model_legacy(&drive, &trunc, [8], Some(&options))
.unwrap_err();
match err {
TbError::Other(message) => {
assert!(message.contains("duplicate vector"), "{message}");
}
other => panic!("unexpected error: {other}"),
}
}
#[test]
fn floquet_effective_order0_matches_h0() {
let model = chain_model();
let drive = FloquetDrive::with_modes(
1.1,
vec![LightMode::new(1, arr1(&[Complex::new(0.23, 0.0)]))],
);
let trunc = FloquetTruncation::new(2, 512);
let options = FloquetEffectiveOptions::new().with_order(0);
let effective = model
.floquet_effective_model_legacy(&drive, &trunc, [32], Some(&options))
.unwrap();
assert_eq!(effective.nsta(), model.nsta());
assert_eq!(effective.hamR, model.hamR);
let k = arr1(&[0.173]);
let from_model = effective.gen_ham(&k, Gauge::Lattice);
let harmonic_cache =
model.floquet_harmonic_cache(&drive, &trunc, 0, 0, &PeierlsFourierMethod::TimeGrid);
let h0 = model.floquet_cached_harmonic_onek(&k, 0, Gauge::Lattice, &harmonic_cache);
let mut max_diff = 0.0f64;
for i in 0..from_model.nrows() {
for j in 0..from_model.ncols() {
max_diff = max_diff.max((from_model[[i, j]] - h0[[i, j]]).norm());
}
}
assert!(max_diff < 1e-12, "order-0 effective mismatch: {max_diff:e}");
}
#[test]
fn floquet_effective_no_drive_matches_static_model() {
let model = chain_model();
let drive = FloquetDrive::new(0.9);
let trunc = FloquetTruncation::new(2, 64);
let effective = model
.floquet_effective_model_legacy(&drive, &trunc, [32], None)
.unwrap();
assert_eq!(effective.nsta(), model.nsta());
assert_eq!(effective.hamR, model.hamR);
let k = arr1(&[0.271]);
for gauge in [Gauge::Lattice, Gauge::Atom] {
let from_effective = effective.gen_ham(&k, gauge);
let from_static = model.gen_ham(&k, gauge);
let mut max_diff = 0.0f64;
for i in 0..from_effective.nrows() {
for j in 0..from_effective.ncols() {
max_diff = max_diff.max((from_effective[[i, j]] - from_static[[i, j]]).norm());
}
}
assert!(
max_diff < 1e-12,
"no-drive effective mismatch in {gauge:?}: {max_diff:e}"
);
}
}
#[test]
fn floquet_weak_drive_matches_first_order_peierls() {
let model = chain_model();
let amp = 1e-5;
let drive = FloquetDrive::with_modes(
1.0,
vec![LightMode::new(1, arr1(&[Complex::new(amp, 0.0)]))],
);
let trunc = FloquetTruncation::new(1, 512);
let k = arr1(&[0.25]);
let hf = model
.floquet_ham_onek(&k, &drive, &trunc, Gauge::Lattice)
.unwrap();
let nsta = model.nsta();
let sector = |n: isize| -> usize { (n + trunc.n_max) as usize };
let h_q1 = hf[[sector(0) * nsta, sector(-1) * nsta]];
let expected = -amp * (TAU * k[0]).sin();
assert!(
(h_q1.re - expected).abs() < 1e-9,
"got {}, expected {}",
h_q1.re,
expected
);
assert!(h_q1.im.abs() < 1e-9, "imag part = {}", h_q1.im);
}
#[test]
fn floquet_incident_basis_public_api_example() {
let lat = array![[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
let orb = array![[0.0, 0.0, 0.0]];
let mut model = Model::<false, 3>::tb_model(lat, orb, None).unwrap();
model.set_hop(-1.0_f64, 0, 0, &arr1(&[1isize, 0, 0]), None);
model.set_hop(-0.8_f64, 0, 0, &arr1(&[0isize, 1, 0]), None);
model.set_hop(-0.6_f64, 0, 0, &arr1(&[0isize, 0, 1]), None);
let incident = IncidentBasis::from_direction(&arr1(&[0.0, 0.0, 1.0])).unwrap();
let circular = incident.polarization([
Complex::new(1.0 / 2.0_f64.sqrt(), 0.0),
Complex::new(0.0, 1.0 / 2.0_f64.sqrt()),
]);
let drive =
FloquetDrive::with_modes(0.8, vec![LightMode::new(1, circular.mapv(|z| 0.12 * z))]);
let trunc = FloquetTruncation::new(1, 128);
let k = arr1(&[0.2, 0.1, 0.0]);
let hf = model
.floquet_ham_onek(&k, &drive, &trunc, Gauge::Lattice)
.unwrap();
assert_eq!(hf.dim(), (3, 3));
let mut max_diff = 0.0f64;
for i in 0..hf.nrows() {
for j in 0..hf.ncols() {
max_diff = max_diff.max((hf[[i, j]] - hf[[j, i]].conj()).norm());
}
}
assert!(max_diff < 1e-11, "max hermiticity error = {max_diff:e}");
let qe = model
.floquet_quasienergy_onek(&k, &drive, &trunc, Gauge::Lattice)
.unwrap();
assert_eq!(qe.len(), 3);
for &x in qe.iter() {
assert!(
x >= -0.5 * drive.omega0_ev - 1e-12 && x < 0.5 * drive.omega0_ev + 1e-12,
"quasienergy {x} is outside the first Floquet zone"
);
}
}
fn two_band_qwz(
tx: f64,
ty: f64,
tz: f64,
m: f64,
lat: [[f64; 2]; 2],
) -> Model<false, 2, NoRMatrix> {
let lat = array![lat[0], lat[1]];
let orb = array![[0.0, 0.0], [0.0, 0.0]];
let mut model = Model::<false, 2>::tb_model(lat, orb, None).unwrap();
model.set_hop(m, 0, 0, &array![0, 0], None);
model.set_hop(-m, 1, 1, &array![0, 0], None);
for r in [[1, 0], [0, 1]] {
model.set_hop(-tz, 0, 0, &array![r[0], r[1]], None);
model.set_hop(tz, 1, 1, &array![r[0], r[1]], None);
}
model.set_hop(tx / 2.0, 0, 1, &array![1, 0], None);
model.set_hop(tx / 2.0, 0, 1, &array![-1, 0], None);
model.set_hop(-ty / 2.0, 0, 1, &array![0, 1], None);
model.set_hop(ty / 2.0, 0, 1, &array![0, -1], None);
model
}
fn qwz_d(k: &[f64; 2], tx: f64, ty: f64, tz: f64, m: f64) -> [f64; 3] {
let (kx, ky) = (TAU * k[0], TAU * k[1]);
[
tx * kx.cos(),
ty * ky.sin(),
m - 2.0 * tz * (kx.cos() + ky.cos()),
]
}
fn decompose_two_band(h: &Array2<Complex<f64>>) -> (f64, [f64; 3]) {
let eps = (h[[0, 0]] + h[[1, 1]]).re / 2.0;
let dz = (h[[0, 0]] - h[[1, 1]]).re / 2.0;
let dx = h[[0, 1]].re;
let dy = -h[[0, 1]].im;
(eps, [dx, dy, dz])
}
fn independent_dressed_harmonic(
model: &Model<false, 2, NoRMatrix>,
k: &[f64; 2],
drive: &FloquetDrive,
q: isize,
n_time: usize,
) -> Array2<Complex<f64>> {
let nsta = model.nsta();
let norb = model.norb();
let mut hq = Array2::<Complex<f64>>::zeros((nsta, nsta));
for it in 0..n_time {
let theta = TAU * (it as f64) / (n_time as f64);
let mut a = [0.0f64; 2];
for mode in &drive.modes {
let phase = Complex::new(0.0, -(mode.harmonic as f64) * theta).exp();
for (comp, ai) in mode.a_complex.iter().enumerate() {
a[comp] += (ai * phase).re;
}
}
let mut h = Array2::<Complex<f64>>::zeros((nsta, nsta));
for (i_r, r_row) in model.hamR.outer_iter().enumerate() {
let r = [r_row[0], r_row[1]];
let bloch =
Complex::new(0.0, TAU * (r[0] as f64 * k[0] + r[1] as f64 * k[1])).exp();
for i in 0..nsta {
for j in 0..nsta {
let t = model.ham[[i_r, i, j]];
if t.norm_sqr() == 0.0 {
continue;
}
let mut d = [0.0f64; 2];
for c in 0..2 {
let mut acc = 0.0;
for b in 0..2 {
let frac = r[b] as f64 + model.orb[[j % norb, b]]
- model.orb[[i % norb, b]];
acc += frac * model.lat[[b, c]];
}
d[c] = acc;
}
let peierls = Complex::new(0.0, -(a[0] * d[0] + a[1] * d[1])).exp();
h[[i, j]] += t * bloch * peierls;
}
}
}
let fourier = Complex::new(0.0, (q as f64) * theta).exp();
hq.scaled_add(fourier, &h);
}
hq.mapv(|x| x / (n_time as f64))
}
fn independent_heff(
model: &Model<false, 2, NoRMatrix>,
k: &[f64; 2],
drive: &FloquetDrive,
q_max: isize,
n_time: usize,
) -> Array2<Complex<f64>> {
let mut h_eff = independent_dressed_harmonic(model, k, drive, 0, n_time);
for q in 1..=q_max {
let hp = independent_dressed_harmonic(model, k, drive, q, n_time);
let hm = independent_dressed_harmonic(model, k, drive, -q, n_time);
let comm = hp.dot(&hm) - hm.dot(&hp);
let scale = Complex::new(1.0 / ((q as f64) * drive.omega0_ev), 0.0);
h_eff.scaled_add(scale, &comm);
}
h_eff
}
fn code_heff(
model: &Model<false, 2, NoRMatrix>,
k: &[f64; 2],
drive: &FloquetDrive,
n_max: isize,
q_max: isize,
) -> Array2<Complex<f64>> {
let trunc = FloquetTruncation::new(n_max, 512);
let options = FloquetEffectiveOptions::new().with_q_max(q_max);
let eff = model
.floquet_effective_model(drive, &trunc, Some(&options))
.unwrap();
eff.gen_ham(&array![k[0], k[1]], Gauge::Lattice)
}
fn code_heff_order0(
model: &Model<false, 2, NoRMatrix>,
k: &[f64; 2],
drive: &FloquetDrive,
n_max: isize,
) -> Array2<Complex<f64>> {
let trunc = FloquetTruncation::new(n_max, 512);
let options = FloquetEffectiveOptions::new().with_order(0);
let eff = model
.floquet_effective_model(drive, &trunc, Some(&options))
.unwrap();
eff.gen_ham(&array![k[0], k[1]], Gauge::Lattice)
}
fn circular_drive(kappa: f64, eta: f64, omega: f64) -> FloquetDrive {
FloquetDrive::with_modes(
omega,
vec![LightMode::new(
1,
array![Complex::new(kappa, 0.0), Complex::new(0.0, eta * kappa)],
)],
)
}
fn elliptical_drive(ax: f64, ay: f64, omega: f64) -> FloquetDrive {
FloquetDrive::with_modes(
omega,
vec![LightMode::new(
1,
array![Complex::new(ax, 0.0), Complex::new(0.0, ay)],
)],
)
}
fn exotic_drive(kappa: f64, alpha: f64, omega: f64) -> FloquetDrive {
FloquetDrive::with_modes(
omega,
vec![
LightMode::new(1, array![Complex::new(kappa, 0.0), Complex::new(0.0, 0.0)]),
LightMode::new(
2,
array![
Complex::new(0.0, 0.0),
Complex::new(kappa * alpha.sin(), kappa * alpha.cos()),
],
),
],
)
}
#[test]
fn two_band_level1_matches_independent_integration() {
let (tx, ty, tz, m) = (1.0, 0.7, 0.5, 0.8);
let omega = 8.0;
let n_time = 4096;
let q_max = 4;
let ks = [[0.13, 0.27], [0.44, 0.61]];
let lattices = [
("square", [[1.0, 0.0], [0.0, 1.0]]),
("rectangular", [[1.0, 0.0], [0.0, 1.6]]),
];
for (lname, lat) in lattices {
let model = two_band_qwz(tx, ty, tz, m, lat);
let drives: Vec<(&str, FloquetDrive)> = vec![
("circular", circular_drive(0.5, 1.0, omega)),
("elliptical", elliptical_drive(0.5, 0.3, omega)),
("exotic", exotic_drive(0.4, 0.6, omega)),
];
for (dname, drive) in drives {
for k in ks {
let code = code_heff(&model, &k, &drive, 2, q_max);
let indep = independent_heff(&model, &k, &drive, q_max, n_time);
for i in 0..2 {
for j in 0..2 {
assert!(
(code[[i, j]] - indep[[i, j]]).norm() < 1e-8,
"[{lname}/{dname}] k={k:?}: code {} vs independent {}",
code[[i, j]],
indep[[i, j]]
);
}
}
}
}
}
}
#[test]
fn two_band_weak_field_matches_cross_product() {
let (tx, ty, tz, m) = (1.0, 0.7, 0.5, 0.8);
let omega = 8.0;
let k = [0.17, 0.31];
let model = two_band_qwz(tx, ty, tz, m, [[1.0, 0.0], [0.0, 1.0]]);
let (kx, ky) = (TAU * k[0], TAU * k[1]);
let dx_d = [-TAU * tx * kx.sin(), 0.0, 2.0 * TAU * tz * kx.sin()];
let dy_d = [0.0, TAU * ty * ky.cos(), 2.0 * TAU * tz * ky.sin()];
let cross = [
dx_d[1] * dy_d[2] - dx_d[2] * dy_d[1],
dx_d[2] * dy_d[0] - dx_d[0] * dy_d[2],
dx_d[0] * dy_d[1] - dx_d[1] * dy_d[0],
];
let lap = [
-TAU * TAU * tx * kx.cos(),
-TAU * TAU * ty * ky.sin(),
2.0 * tz * TAU * TAU * (kx.cos() + ky.cos()),
];
let d_static = qwz_d(&k, tx, ty, tz, m);
let d_eff = |kappa: f64, eta: f64| -> [f64; 3] {
let drive = circular_drive(kappa, eta, omega);
let h = code_heff(&model, &k, &drive, 1, 1);
decompose_two_band(&h).1
};
let cpl = |kappa: f64| -> [f64; 3] {
let dp = d_eff(kappa, 1.0);
let dm = d_eff(kappa, -1.0);
[
(dp[0] - dm[0]) / 2.0,
(dp[1] - dm[1]) / 2.0,
(dp[2] - dm[2]) / 2.0,
]
};
let a2 = |kappa: f64| -> [f64; 3] {
let dp = d_eff(kappa, 1.0);
let dm = d_eff(kappa, -1.0);
[
(dp[0] + dm[0]) / 2.0 - d_static[0],
(dp[1] + dm[1]) / 2.0 - d_static[1],
(dp[2] + dm[2]) / 2.0 - d_static[2],
]
};
let richardson = |f1: [f64; 3], f2: [f64; 3], k1: f64| -> [f64; 3] {
[
(16.0 * f2[0] - f1[0]) / (3.0 * k1 * k1),
(16.0 * f2[1] - f1[1]) / (3.0 * k1 * k1),
(16.0 * f2[2] - f1[2]) / (3.0 * k1 * k1),
]
};
let kappa1 = 0.1;
let cpl_coeff = richardson(cpl(kappa1), cpl(kappa1 / 2.0), kappa1);
let a2_coeff = richardson(a2(kappa1), a2(kappa1 / 2.0), kappa1);
let cpl_pred = [
cross[0] / (TAU * TAU * omega),
cross[1] / (TAU * TAU * omega),
cross[2] / (TAU * TAU * omega),
];
let a2_pred = [
lap[0] / (4.0 * TAU * TAU),
lap[1] / (4.0 * TAU * TAU),
lap[2] / (4.0 * TAU * TAU),
];
for comp in 0..3 {
assert!(
(cpl_coeff[comp] - cpl_pred[comp]).abs() < 1e-4,
"CPL coefficient[{comp}]: {:.6} vs analytic {:.6}",
cpl_coeff[comp],
cpl_pred[comp]
);
assert!(
(a2_coeff[comp] - a2_pred[comp]).abs() < 1e-4,
"A² coefficient[{comp}]: {:.6} vs analytic {:.6}",
a2_coeff[comp],
a2_pred[comp]
);
}
let d_eff_ell = |ax: f64, ay: f64| -> [f64; 3] {
let drive = elliptical_drive(ax, ay, omega);
let h = code_heff(&model, &k, &drive, 1, 1);
decompose_two_band(&h).1
};
let (ax, ay) = (0.08, 0.05);
let dp = d_eff_ell(ax, ay);
let dm = d_eff_ell(ax, -ay);
let cpl_ell = [
(dp[0] - dm[0]) / 2.0,
(dp[1] - dm[1]) / 2.0,
(dp[2] - dm[2]) / 2.0,
];
for comp in 0..3 {
let pred = ax * ay * cross[comp] / (TAU * TAU * omega);
assert!(
(cpl_ell[comp] - pred).abs() < 1e-5,
"elliptical CPL[{comp}]: {:.6} vs analytic {:.6}",
cpl_ell[comp],
pred
);
}
}
#[test]
fn two_band_exotic_drive_first_order_commutator_is_cubic() {
let (tx, ty, tz, m) = (1.0, 0.7, 0.5, 0.8);
let omega = 8.0;
let k = [0.21, 0.37];
let model = two_band_qwz(tx, ty, tz, m, [[1.0, 0.0], [0.0, 1.0]]);
let commutator_norm = |_kappa: f64, drive: FloquetDrive| -> f64 {
let h1 = code_heff(&model, &k, &drive, 2, 4);
let h0 = code_heff_order0(&model, &k, &drive, 2);
let d1 = decompose_two_band(&h1).1;
let d0 = decompose_two_band(&h0).1;
((d1[0] - d0[0]).powi(2) + (d1[1] - d0[1]).powi(2) + (d1[2] - d0[2]).powi(2)).sqrt()
};
let exo1 = commutator_norm(0.3, exotic_drive(0.3, 0.6, omega));
let exo2 = commutator_norm(0.15, exotic_drive(0.15, 0.6, omega));
let ratio_exo = exo1 / exo2;
assert!(
exo1 > 1e-8,
"exotic-drive commutator must be non-vanishing at O(κ³), got {exo1:e}"
);
assert!(
(ratio_exo - 8.0).abs() < 1.0,
"exotic-drive commutator should scale as κ³ (ratio ≈ 8), got {ratio_exo:.2}"
);
let circ1 = commutator_norm(0.3, circular_drive(0.3, 1.0, omega));
let circ2 = commutator_norm(0.15, circular_drive(0.15, 1.0, omega));
let ratio_circ = circ1 / circ2;
assert!(
(ratio_circ - 4.0).abs() < 0.5,
"circular-drive commutator should scale as κ² (ratio ≈ 4), got {ratio_circ:.2}"
);
}
}