use ndarray::prelude::*;
use ndarray_linalg::*;
use num_complex::Complex;
use crate::Gauge;
use crate::Model;
use crate::RMatrixData;
use crate::SpinDirection;
use crate::math::anti_comm;
use super::helpers::build_spin_matrix;
use super::types::VertexKernel;
impl<const SPIN: bool, const DIM: usize, R: RMatrixData> Model<SPIN, DIM, R> {
pub(crate) fn compute_velocity_kernel(
&self,
k_vec: &Array1<f64>,
dir_a: &Array1<f64>,
dir_b: &Array1<f64>,
dir_c: Option<&Array1<f64>>,
gauge: Gauge,
spin: Option<SpinDirection>,
) -> VertexKernel {
let nsta = self.nsta();
let n_dir = if dir_c.is_some() { 3 } else { 2 };
let mut directions = Array2::<f64>::zeros((n_dir, DIM));
directions.row_mut(0).assign(dir_a);
directions.row_mut(1).assign(dir_b);
if let Some(dc) = dir_c {
directions.row_mut(2).assign(dc);
}
let (v_proj, hamk) = self.gen_v_projected(k_vec, gauge, &directions);
let (band, evec) = hamk.eigh(UPLO::Lower).unwrap();
let ut = evec.t();
let uc = evec.map(|x| x.conj());
let to_band = |d: usize, spin_dress: bool| -> Array2<Complex<f64>> {
let v_raw = v_proj.slice(s![d, .., ..]).to_owned();
if spin_dress && SPIN && spin.is_some() {
let x = build_spin_matrix(self.norb(), spin);
let s = anti_comm(&x, &v_raw) * 0.5;
ut.dot(&s.dot(&uc))
} else {
ut.dot(&v_raw.dot(&uc))
}
};
let va = to_band(0, true);
let vb = to_band(1, false);
let mut k_ab = Array2::<Complex<f64>>::zeros((nsta, nsta));
for n in 0..nsta {
for m in 0..nsta {
k_ab[[n, m]] = va[[n, m]] * vb[[m, n]];
}
}
let (vdiag, k_bc, k_ac, vdiag_a, vdiag_b) = if let Some(_dc) = dir_c {
let vc = to_band(2, false);
let mut bc = Array2::<Complex<f64>>::zeros((nsta, nsta));
let mut ac = Array2::<Complex<f64>>::zeros((nsta, nsta));
for n in 0..nsta {
for m in 0..nsta {
bc[[n, m]] = vb[[n, m]] * vc[[m, n]];
ac[[n, m]] = va[[n, m]] * vc[[m, n]];
}
}
(
Some(vc.diag().map(|x| x.re).to_owned()),
Some(bc),
Some(ac),
Some(va.diag().map(|x| x.re).to_owned()),
Some(vb.diag().map(|x| x.re).to_owned()),
)
} else {
(None, None, None, None, None)
};
VertexKernel {
band,
evec,
k_ab,
k_bc,
k_ac,
vdiag,
vdiag_a,
vdiag_b,
}
}
}