use crate::Gauge;
use crate::Model;
use crate::comm;
use crate::solve_ham::*;
use ndarray::prelude::*;
use ndarray::*;
use ndarray_linalg::conjugate;
use ndarray_linalg::*;
use num_complex::Complex;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::f64::consts::PI;
use std::ops::AddAssign;
pub trait Velocity {
fn gen_v<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
gauge: Gauge,
) -> (Array3<Complex<f64>>, Array2<Complex<f64>>);
}
impl Velocity for Model {
#[allow(non_snake_case)]
#[inline(always)]
fn gen_v<S: Data<Elem = f64>>(
&self,
kvec: &ArrayBase<S, Ix1>,
gauge: Gauge,
) -> (Array3<Complex<f64>>, Array2<Complex<f64>>) {
assert_eq!(
kvec.len(),
self.dim_r(),
"Wrong, the k-vector's length {} must equal to the dimension of model {}.",
kvec.len(),
self.dim_r()
);
let Us = (self.hamR.mapv(|x| x as f64))
.dot(kvec)
.mapv(|x| Complex::<f64>::new(x, 0.0));
let Us = Us * Complex::new(0.0, 2.0 * PI);
let Us = Us.mapv(Complex::exp); let mut v = Array3::<Complex<f64>>::zeros((self.dim_r(), self.nsta(), self.nsta()));
let R0 = &self.hamR.mapv(|x| Complex::<f64>::new(x as f64, 0.0));
let R0 = R0.dot(&self.lat.mapv(|x| Complex::new(x, 0.0)));
let hamk: Array2<Complex<f64>> = self
.ham
.outer_iter()
.zip(Us.iter())
.fold(Array2::zeros((self.nsta(), self.nsta())), |acc, (hm, u)| {
acc + &hm * *u
});
let (v, hamk) = match gauge {
Gauge::Atom => {
let orb_sta = if self.spin {
let orb0 = concatenate(Axis(0), &[self.orb.view(), self.orb.view()]).unwrap();
orb0
} else {
self.orb.to_owned()
};
let U0 = orb_sta.dot(kvec);
let U0 = U0.mapv(|x| Complex::<f64>::new(x, 0.0));
let U0 = U0 * Complex::new(0.0, 2.0 * PI);
let mut U0 = U0.mapv(Complex::exp);
let U = Array2::from_diag(&U0);
let U_conj = Array2::from_diag(&U0.mapv(|x| x.conj()));
let orb_real = orb_sta.dot(&self.lat);
let mut UU = Array3::<f64>::zeros((self.dim_r(), self.nsta(), self.nsta()));
let A = orb_real.view().insert_axis(Axis(2));
let A = A
.broadcast((self.nsta(), self.dim_r(), self.nsta()))
.unwrap()
.permuted_axes([1, 0, 2]);
let mut B = A.view().permuted_axes([0, 2, 1]);
let UU = &B - &A;
let UU = UU.mapv(|x| Complex::<f64>::new(0.0, x)); Zip::from(v.outer_iter_mut())
.and(R0.axis_iter(Axis(1)))
.and(UU.outer_iter())
.for_each(|mut v0, r, det_tau| {
let vv: Array2<Complex<f64>> =
self.ham.outer_iter().zip(Us.iter().zip(r.iter())).fold(
Array2::zeros((self.nsta(), self.nsta())),
|acc, (ham, (us, r0))| acc + &ham * *us * *r0 * Complex::i(),
);
let vv: Array2<Complex<f64>> = &vv + &hamk * &det_tau;
let vv = &U_conj.dot(&vv);
let vv = vv.dot(&U); v0.assign(&vv);
});
let hamk = U_conj.dot(&hamk.dot(&U)); if self.rmatrix.len_of(Axis(0)) != 1 {
let mut rk =
Array3::<Complex<f64>>::zeros((self.dim_r(), self.nsta(), self.nsta()));
let mut rk = self
.rmatrix
.axis_iter(Axis(0))
.zip(Us.iter())
.fold(rk, |acc, (ham, us)| acc + &ham * *us);
for i in 0..3 {
let mut r0: ArrayViewMut2<Complex<f64>> = rk.slice_mut(s![i, .., ..]);
let r_new = r0.dot(&U);
let r_new = U_conj.dot(&r_new);
r0.assign(&r_new);
let mut dig = r0.diag_mut();
dig.assign(&Array1::zeros(self.nsta()));
let A = comm(&hamk, &r0) * Complex::i();
v.slice_mut(s![i, .., ..]).add_assign(&A);
}
}
(v, hamk)
}
Gauge::Lattice => {
Zip::from(v.outer_iter_mut())
.and(R0.axis_iter(Axis(1)))
.for_each(|mut v0, r| {
let vv: Array2<Complex<f64>> =
self.ham.outer_iter().zip(Us.iter().zip(r.iter())).fold(
Array2::zeros((self.nsta(), self.nsta())),
|acc, (ham, (us, r0))| acc + &ham * *us * *r0 * Complex::i(),
);
v0.assign(&vv);
});
if self.rmatrix.len_of(Axis(0)) != 1 {
let mut rk =
Array3::<Complex<f64>>::zeros((self.dim_r(), self.nsta(), self.nsta()));
let mut rk = self
.rmatrix
.axis_iter(Axis(0))
.zip(Us.iter())
.fold(rk, |acc, (ham, us)| acc + &ham * *us);
for i in 0..3 {
let mut r0: ArrayViewMut2<Complex<f64>> = rk.slice_mut(s![i, .., ..]);
let A = comm(&hamk, &r0) * Complex::i();
v.slice_mut(s![i, .., ..]).add_assign(&A);
}
}
(v, hamk)
}
};
(v, hamk)
}
}