Rustb/

lib.rs

#![allow(warnings)]
pub mod SKmodel;
pub mod atom_struct;
pub mod conductivity;
pub mod cut;
pub mod error;
pub mod generics;
pub mod geometry;
pub mod io;
pub mod kpath;
pub mod kpoints;
pub mod magnetic_field;
pub mod math;
pub mod model;
pub mod model_build;
pub mod model_physics;
pub mod model_utils;
pub mod ndarray_lapack;
pub mod optical_conductivity;
pub mod orbital_angular;
pub mod output;
pub mod phy_const;
pub mod solve_ham;
pub mod surfgreen;
pub mod unfold;
pub mod velocity;
pub mod wannier90;
pub use crate::SKmodel::{SkAtom, SkParams, SlaterKosterModel, ToTbModel};
pub use crate::atom_struct::{Atom, OrbProj};
pub use crate::conductivity::*;
pub use crate::cut::CutModel;
pub use crate::error::{Result, TbError};
use crate::generics::usefloat;
pub use crate::geometry::*;
pub use crate::io::*;
pub use crate::kpath::*;
pub use crate::kpoints::{gen_kmesh, gen_krange};
pub use crate::magnetic_field::*;
pub use crate::math::*;
pub use crate::model::*;
pub use crate::model_physics::*;
pub use crate::optical_conductivity::*;
pub use crate::output::*;
pub use crate::solve_ham::solve;
pub use crate::surfgreen::surf_Green;
pub use crate::unfold::Unfold;
pub use crate::velocity::*;
pub use crate::wannier90::*;

/// # Rustb - Tight-Binding Model Library
///
/// A comprehensive Rust library for tight-binding model calculations with support for:
/// - Band structure calculations
/// - Surface state computations using Green's functions
/// - Linear and nonlinear transport properties
/// - Berry phase and curvature calculations
/// - Slater-Koster parameterized models
///
/// ## Key Features
///
/// - **Band Structure**: Solve eigenvalue problems $H(\mathbf{k}) \psi_n(\mathbf{k}) = E_n(\mathbf{k}) \psi_n(\mathbf{k})$
/// - **Surface States**: Compute surface Green's functions $G^s(\omega, \mathbf{k}_\parallel)$
/// - **Transport Properties**:
///   - Anomalous Hall conductivity $\sigma_{xy} = \frac{e^2}{\hbar} \int \frac{d^2k}{(2\pi)^2} \Omega_z(\mathbf{k})$
///   - Spin Hall conductivity
///   - Nonlinear conductivity tensors
/// - **Topological Invariants**: Chern numbers, Wilson loops, and Wannier centers
/// - **Slater-Koster Models**: Parameterized tight-binding models with two-center integrals
/// - **File I/O**: Utilities for writing calculation results to text files
///
/// ## Mathematical Foundation
///
/// The library implements the tight-binding Hamiltonian:
/// $$
/// H = \sum_{i,j} t_{ij} c_i^\dagger c_j + \sum_i \epsilon_i c_i^\dagger c_i
/// $$
/// where $t_{ij}$ are hopping parameters and $\epsilon_i$ are on-site energies.
///
/// For transport calculations, we compute the Berry curvature:
/// $$
/// \Omega_n(\mathbf{k}) = -2\,\text{Im}\sum_{m\neq n} \frac{\bra{n}\partial_{k_x} H\ket{m}\bra{m}\partial_{k_y} H\ket{n}}{(E_n - E_m)^2}
/// $$
///
/// ## File I/O Utilities
///
/// The library provides file output utilities through the `io` module:
/// - `write_txt`: Write 2D arrays to text files with formatted output
/// - `write_txt_1`: Write 1D arrays to text files with formatted output
///
/// These functions automatically handle number formatting and spacing for scientific data.

///An example
///
///```
///use gnuplot::{Color,Figure, AxesCommon, AutoOption::Fix,HOT};
///use gnuplot::Major;
///use ndarray::*;
///use ndarray::prelude::*;
///use num_complex::Complex;
///use Rustb::*;
///
///fn graphene(){
///    let li:Complex<f64>=1.0*Complex::i();
///    let t1=1.0+0.0*li;
///    let t2=0.1+0.0*li;
///    let t3=0.0+0.0*li;
///    let delta=0.5;
///    let dim_r:usize=2;
///    let norb:usize=2;
///    let lat=arr2(&[[3.0_f64.sqrt(),-1.0],[3.0_f64.sqrt(),1.0]]);
///    let orb=arr2(&[[0.0,0.0],[1.0/3.0,1.0/3.0]]);
///    let mut model=Model::tb_model(dim_r,lat,orb,false,None);
///    model.set_onsite(&arr1(&[delta,-delta]),SpinDirection::None);
///    model.add_hop(t1,0,1,&array![0,0],SpinDirection::None);
///    model.add_hop(t1,0,1,&array![-1,0],SpinDirection::None);
///    model.add_hop(t1,0,1,&array![0,-1],SpinDirection::None);
///    model.add_hop(t2,0,0,&array![1,0],SpinDirection::None);
///    model.add_hop(t2,1,1,&array![1,0],SpinDirection::None);
///    model.add_hop(t2,0,0,&array![0,1],SpinDirection::None);
///    model.add_hop(t2,1,1,&array![0,1],SpinDirection::None);
///    model.add_hop(t2,0,0,&array![1,-1],SpinDirection::None);
///    model.add_hop(t2,1,1,&array![1,-1],SpinDirection::None);
///    model.add_hop(t3,0,1,&array![1,-1],SpinDirection::None);
///    model.add_hop(t3,0,1,&array![-1,1],SpinDirection::None);
///    model.add_hop(t3,0,1,&array![-1,-1],SpinDirection::None);
///    let nk:usize=1001;
///    let path=[[0.0,0.0],[2.0/3.0,1.0/3.0],[0.5,0.5],[0.0,0.0]];
///    let path=arr2(&path);
///    let (k_vec,k_dist,k_node)=model.k_path(&path,nk);
///    let (eval,evec)=model.solve_all_parallel(&k_vec);
///    let label=vec!["G","K","M","G"];
///    let (k_vec,k_dist,k_node)=model.k_path(&path,nk); //generate the k vector
///    let eval=model.solve_band_all_parallel(&k_vec);  //calculate the bands
///    let mut fg = Figure::new();
///    let x:Vec<f64>=k_dist.to_vec();
///    let axes=fg.axes2d();
///    for i in 0..model.nsta{
///        let y:Vec<f64>=eval.slice(s![..,i]).to_owned().to_vec();
///        axes.lines(&x, &y, &[Color("black")]);
///    }
///    let axes=axes.set_x_range(Fix(0.0), Fix(k_node[[k_node.len()-1]]));
///    let label=label.clone();
///    let mut show_ticks=Vec::new();
///    for i in 0..k_node.len(){
///        let A=k_node[[i]];
///        let B=label[i];
///        show_ticks.push(Major(A,Fix(B)));
///    }
///    axes.set_x_ticks_custom(show_ticks.into_iter(),&[],&[]);
///    let k_node=k_node.to_vec();
///    let mut jpg_name=String::new();
///    jpg_name.push_str("band.jpg");
///    fg.set_terminal("jpeg", &jpg_name);
///    fg.show();
///
///    //start to draw the band structure
///    //Starting to calculate the edge state, first is the zigzag state
///    let nk:usize=501;
///    let U=arr2(&[[1.0,1.0],[-1.0,1.0]]);
///    let super_model=model.make_supercell(&U);
///    let zig_model=super_model.cut_piece(100,0);
///    let path=[[0.0,0.0],[0.0,0.5],[0.0,1.0]];
///    let path=arr2(&path);
///    let label=vec!["G","M","G"];
///    zig_model.show_band(&path,&label,nk,"graphene_zig");
///    //Starting to calculate the DOS of graphene
///    let nk:usize=101;
///    let kmesh=arr1(&[nk,nk]);
///    let E_min=-3.0;
///    let E_max=3.0;
///    let E_n=1000;
///    let (E0,dos)=model.dos(&kmesh,E_min,E_max,E_n,1e-2);
///    //start to show DOS
///    let mut fg = Figure::new();
///    let x:Vec<f64>=E0.to_vec();
///    let axes=fg.axes2d();
///    let y:Vec<f64>=dos.to_vec();
///    axes.lines(&x, &y, &[Color("black")]);
///    let mut show_ticks=Vec::<String>::new();
///    let mut pdf_name=String::new();
///    pdf_name.push_str("dos.jpg");
///    fg.set_terminal("pdfcairo", &pdf_name);
///    fg.show();
///}
///```
///
///

#[cfg(test)]
mod tests {
    use super::*;
    use gnuplot::{
        AutoOption, AxesCommon, Color, Figure, Fix, Font, LineStyle, Major, PointSymbol, Rotate,
        Solid, TextOffset,
    };
    use ndarray::concatenate;
    use ndarray::linalg::kron;
    use ndarray::prelude::*;
    use ndarray::*;
    use ndarray_linalg::conjugate;
    use ndarray_linalg::*;
    use ndarray_linalg::{Eigh, UPLO};
    use num_complex::Complex;
    use rayon::prelude::*;
    use std::f64::consts::PI;
    use std::fs::File;
    use std::io::Write;
    use std::time::{Duration, Instant};
    use std::fs::create_dir_all;

    fn write_txt(data: Array2<f64>, output: &str) -> std::io::Result<()> {
        let mut file = File::create(output).expect("Unable to BAND.dat");
        let n = data.len_of(Axis(0));
        let s = data.len_of(Axis(1));
        let mut s0 = String::new();
        for i in 0..n {
            for j in 0..s {
                if data[[i, j]] >= 0.0 {
                    s0.push_str("     ");
                } else {
                    s0.push_str("    ");
                }
                let aa = format!("{:.6}", data[[i, j]]);
                s0.push_str(&aa);
            }
            s0.push_str("\n");
        }
        writeln!(file, "{}", s0)?;
        Ok(())
    }

    fn write_txt_1(data: Array1<f64>, output: &str) -> std::io::Result<()> {
        use std::fs::File;
        use std::io::Write;
        let mut file = File::create(output).expect("Unable to BAND.dat");
        let n = data.len_of(Axis(0));
        let mut s0 = String::new();
        for i in 0..n {
            if data[[i]] >= 0.0 {
                s0.push_str(" ");
            }
            let aa = format!("{:.6}\n", data[[i]]);
            s0.push_str(&aa);
        }
        writeln!(file, "{}", s0)?;
        Ok(())
    }
    #[test]
    fn test_gen_v() {
        //判断两个Array1<f64> 是否足够接近
        fn are_arrays_close(a: &Array1<f64>, b: &Array1<f64>, tolerance: f64) -> bool {
            a.iter()
                .zip(b.iter())
                .all(|(&x, &y)| (x - y).abs() < tolerance)
        }

        //判断两个Array2<Compelx<f64>> 是否足够接近
        fn are_complex_arrays_close(
            a: &Array2<Complex<f64>>,
            b: &Array2<Complex<f64>>,
            tolerance: f64,
        ) -> bool {
            a.iter()
                .zip(b.iter())
                .all(|(&x, &y)| (x.re - y.re).abs() < tolerance && (x.im - y.im).abs() < tolerance)
        }
        let li: Complex<f64> = 1.0 * Complex::i();
        let t = 1.0;
        let delta = 0.0;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[1.0, 0.0], [0.5, 3.0_f64.sqrt() / 2.0]]);
        let orb = arr2(&[[1.0 / 3.0, 1.0 / 3.0], [2.0 / 3.0, 2.0 / 3.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[-delta, delta]), SpinDirection::None);
        let R0: Array2<isize> = arr2(&[[0, 0], [-1, 0], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.set_hop(t, 0, 1, &R, SpinDirection::None);
        }
        assert_eq!(model.solve_band_onek(&array![0.0, 0.0]), array![-3.0, 3.0]);
        let result = model.solve_band_onek(&array![1.0 / 3.0, 2.0 / 3.0]);
        assert!(
            are_arrays_close(&result, &array![0.0, 0.0], 1e-5),
            "wrong!, the solve_band_onek get wrong result! please check it!"
        );
        let (result, _) = model.gen_v(&array![1.0 / 3.0, 1.0 / 3.0], Gauge::Atom);
        let resulty = array![
            [0.0 * li, -0.4698463103929542 - 0.17101007166283436 * li],
            [-0.4698463103929542 + 0.17101007166283436 * li, 0.0 * li]
        ];
        let resultx = array![
            [0.0 * li, -0.8137976813493737 - 0.2961981327260237 * li],
            [-0.8137976813493737 + 0.2961981327260237 * li, 0.0 * li]
        ];
        println!("result={}", result);
        assert!(
            are_complex_arrays_close(&result.slice(s![0, .., ..]).to_owned(), &resultx, 1e-8),
            "Wrong! the gen_v is get wrong results! please check it!"
        );
        assert!(
            are_complex_arrays_close(&result.slice(s![1, .., ..]).to_owned(), &resulty, 1e-8),
            "Wrong! the gen_v is get wrong results! please check it!"
        );

        let (result, _) = model.gen_v(&array![1.0 / 3.0, 1.0 / 3.0], Gauge::Lattice);
        let resultx = array![
            [
                0.0 * li,
                -3.0 * 3.0_f64.sqrt() / 4.0 * t + 3.0 / 4.0 * t * li
            ],
            [
                -3.0 * 3.0_f64.sqrt() / 4.0 * t - 3.0 / 4.0 * t * li,
                0.0 * li
            ]
        ];
        println!("result={}", &result - &resultx);
        assert!(
            are_complex_arrays_close(&result.slice(s![0, .., ..]).to_owned(), &resultx, 1e-8),
            "Wrong! the gen_v is get wrong results! please check it!"
        );

        let kvec = array![1.0 / 3.0, 1.0 / 3.0];
        let (band, evec) = model.solve_onek(&kvec);
        let ham = model.gen_ham(&kvec, Gauge::Atom);
        let evec_conj = evec.map(|x| x.conj());
        let evec = evec.t();
        let ham = ham.dot(&evec);
        let ham = evec_conj.dot(&ham);
        let new_band = ham.diag().map(|x| x.re);
        assert!(
            are_arrays_close(&new_band, &band, 1e-5),
            "wrong!, the solve_onek get wrong result! please check it!"
        );
    }
    #[test]
    fn conductivity_test() {
        //这个是用 Haldan 模型来测试
        let li: Complex<f64> = 1.0 * Complex::i();
        let t = -1.0 + 0.0 * li;
        let t2 = -1.0 + 0.0 * li;
        let delta = 0.7;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[1.0, 0.0], [0.5, 3.0_f64.sqrt() / 2.0]]);
        let orb = arr2(&[[1.0 / 3.0, 1.0 / 3.0], [2.0 / 3.0, 2.0 / 3.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[-delta, delta]), SpinDirection::None);
        let R0: Array2<isize> = arr2(&[[0, 0], [-1, 0], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t, 0, 1, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[1, 0], [-1, 1], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t2 * li, 0, 0, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[-1, 0], [1, -1], [0, 1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t2 * li, 1, 1, &R, SpinDirection::None);
        }
        let k_vec = array![1.0 / 3.0, 2.0 / 3.0];
        let dir_1 = array![1.0, 0.0];
        let dir_2 = array![0.0, 1.0];
        let mu = 0.0;
        let T = 0.0;
        let og = 0.0;
        let spin = 0;
        let eta = 1e-3;
        let result1 =
            model.berry_curvature_onek(&k_vec, &dir_1, &dir_2, mu, T, spin, eta) * (2.0 * PI);

        let mut k_list = Array2::zeros((9, 2));
        let dk = 0.0001;
        k_list.row_mut(0).assign(&(&k_vec + dk * &dir_1));
        k_list
            .row_mut(1)
            .assign(&(&k_vec + dk * &dir_1 + dk * &dir_2));
        k_list.row_mut(2).assign(&(&k_vec + dk * &dir_2));
        k_list
            .row_mut(3)
            .assign(&(&k_vec - dk * &dir_1 + dk * &dir_2));
        k_list.row_mut(4).assign(&(&k_vec - dk * &dir_1));
        k_list
            .row_mut(5)
            .assign(&(&k_vec - dk * &dir_1 - dk * &dir_2));
        k_list.row_mut(6).assign(&(&k_vec - dk * &dir_2));
        k_list
            .row_mut(7)
            .assign(&(&k_vec + dk * &dir_1 - dk * &dir_2));
        k_list.row_mut(8).assign(&(&k_vec + dk * &dir_1));
        let result2 = model.berry_loop(&k_list, &vec![0]);
        let result2 = result2[[0]] / (dk.powi(2)) / 4.0 / (2.0 * PI) * 3_f64.sqrt() / 2.0;
        println!("result2={},result1={}", result2, result1);
        assert!(
            (result2 - result1).abs() < 1e-4,
            "Wrong!, the berry_curvature or berry_flux mut be false"
        );
    }
    #[test]
    fn gen_v_speed_test() {
        println!("开始测试各个函数的运行速度, 用次近邻的石墨烯模型");
        let li: Complex<f64> = 1.0 * Complex::i();
        let t = 2.0 + 0.0 * li;
        let t2 = -1.0 + 0.0 * li;
        let delta = 0.7;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[1.0, 0.0], [0.5, 3.0_f64.sqrt() / 2.0]]);
        let orb = arr2(&[[1.0 / 3.0, 1.0 / 3.0], [2.0 / 3.0, 2.0 / 3.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[-delta, delta]), SpinDirection::None);
        let R0: Array2<isize> = arr2(&[[0, 0], [-1, 0], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t, 0, 1, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[1, 0], [-1, 1], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t2 * li, 0, 0, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[-1, 0], [1, -1], [0, 1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t2 * li, 1, 1, &R, SpinDirection::None);
        }
        println!("{:?}", model.atom_list());
        let U = array![[3.0, 0.0], [0.0, 3.0]];
        let model = model.make_supercell(&U).unwrap();

        let nk = 101;
        let k_mesh = array![nk, nk];
        let kvec = gen_kmesh(&k_mesh).unwrap();

        {
            println!("开始计算 gen_v 的耗时速度, 为了平均, 我们单线程求解gen_v");
            let start = Instant::now(); // 开始计时
            let A: Vec<_> = kvec
                .outer_iter()
                .into_par_iter()
                .map(|x| {
                    let (a, _) = model.gen_v(&x.to_owned(), Gauge::Atom);
                    a
                })
                .collect();
            let end = Instant::now(); // 结束计时
            let duration = end.duration_since(start); // 计算执行时间
            println!(
                "run gen_v {} times took {} seconds",
                kvec.nrows(),
                duration.as_secs_f64()
            ); // 输出执行时间
        }
    }
    #[test]
    fn Haldan_model() {
        let li: Complex<f64> = 1.0 * Complex::i();
        let t = -1.0 + 0.0 * li;
        let t2 = -1.0 + 0.0 * li;
        let delta = 0.7;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[1.0, 0.0], [0.5, 3.0_f64.sqrt() / 2.0]]);
        let orb = arr2(&[[1.0 / 3.0, 1.0 / 3.0], [2.0 / 3.0, 2.0 / 3.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[-delta, delta]), SpinDirection::None);
        let R0: Array2<isize> = arr2(&[[0, 0], [-1, 0], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t, 0, 1, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[1, 0], [-1, 1], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t2 * li, 0, 0, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[-1, 0], [1, -1], [0, 1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.add_hop(t2 * li, 1, 1, &R, SpinDirection::None);
        }
        let nk: usize = 101;
        let path = [
            [0.0, 0.0],
            [2.0 / 3.0, 1.0 / 3.0],
            [0.5, 0.5],
            [1.0 / 3.0, 2.0 / 3.0],
            [0.0, 0.0],
        ];
        let path = arr2(&path);
        let (k_vec, k_dist, k_node) = model.k_path(&path, nk).unwrap();
        let (eval, evec) = model.solve_all_parallel(&k_vec);
        let label = vec!["G", "K", "M", "K'", "G"];
        model.show_band(&path, &label, nk, "tests/Haldan").unwrap();
        /////开始计算体系的霍尔电导率//////
        let nk: usize = 31;
        let T: f64 = 0.0;
        let eta: f64 = 0.001;
        let og: f64 = 0.0;
        let mu: f64 = 0.0;
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let dir_3 = arr1(&[0.0, 1.0]);
        let spin: usize = 0;
        let kmesh = arr1(&[nk, nk]);

        let start = Instant::now(); // 开始计时
        let conductivity = model
            .Hall_conductivity(&kmesh, &dir_1, &dir_2, mu, T, spin, eta)
            .unwrap();
        let end = Instant::now(); // 结束计时
        let duration = end.duration_since(start); // 计算执行时间
        println!("quantom_Hall_effect={}", conductivity * (2.0 * PI));
        assert!(
            (conductivity * (2.0 * PI) - 1.0).abs() < 1e-3,
            "Wrong!, the Hall conductivity is wrong!"
        );
        println!("function_a took {} seconds", duration.as_secs_f64()); // 输出执行时间

        let mu = Array1::linspace(-2.0, 2.0, 101);
        let start = Instant::now(); // 开始计时
        let conductivity_mu = model
            .Hall_conductivity_mu(&kmesh, &dir_1, &dir_2, &mu, T, spin, eta)
            .unwrap();
        let end = Instant::now(); // 结束计时
        let duration = end.duration_since(start); // 计算执行时间
        println!("quantom_Hall_effect={}", conductivity_mu[[50]] * (2.0 * PI));
        assert!(
            (conductivity_mu[[50]] - conductivity).abs() < 1e-3,
            "Wrong!, the Hall conductivity is wrong!, Hall_mu's result is {}, but Hall conductivity is {}",
            conductivity_mu[[50]],
            conductivity
        );
        println!("function_a took {} seconds", duration.as_secs_f64()); // 输出执行时间
        let conductivity = model
            .Hall_conductivity(&kmesh, &dir_1, &dir_2, -2.0, T, spin, eta)
            .unwrap();
        assert!(
            (conductivity_mu[[0]] - conductivity).abs() < 1e-3,
            "Wrong!, the Hall conductivity is wrong!, Hall_mu's result is {}, but Hall conductivity is {}",
            conductivity_mu[[0]],
            conductivity
        );
        let conductivity = model
            .Hall_conductivity(&kmesh, &dir_1, &dir_2, 2.0, T, spin, eta)
            .unwrap();
        assert!(
            (conductivity_mu[[100]] - conductivity).abs() < 1e-3,
            "Wrong!, the Hall conductivity is wrong!, Hall_mu's result is {}, but Hall conductivity is {}",
            conductivity_mu[[100]],
            conductivity
        );
        //开始绘图
        let mut fg = Figure::new();
        let x: Vec<f64> = mu.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = (conductivity_mu * 2.0 * PI).to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/Haldan");
        pdf_name.push_str("/hall_mu.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        let mu = 0.0;
        let nk: usize = 31;
        let kmesh = arr1(&[nk, nk]);
        let start = Instant::now(); // 开始计时
        let conductivity = model
            .Hall_conductivity_adapted(&kmesh, &dir_1, &dir_2, mu, T, spin, eta, 0.01, 0.0001)
            .unwrap();
        let end = Instant::now(); // 结束计时
        let duration = end.duration_since(start); // 计算执行时间
        println!("霍尔电导率{}", conductivity * (2.0 * PI));
        assert!(
            (conductivity * (2.0 * PI) - 1.0).abs() < 1e-3,
            "Wrong!, the Hall conductivity is wrong!"
        );
        println!("function_a took {} seconds", duration.as_secs_f64()); // 输出执行时间
        //画一下3000k的时候的费米导数分布
        let T = 100.0;
        let nk: usize = 101;
        let kmesh = arr1(&[nk, nk]);
        println!("{}", kmesh);
        let E_min = -3.0;
        let E_max = 3.0;
        let E_n = 1000;
        let mu = Array1::linspace(E_min, E_max, E_n);
        let beta: f64 = 1.0 / T / (8.617e-5);
        let f: Array1<f64> = 1.0 / ((beta * &mu).mapv(f64::exp) + 1.0);
        let par_f = beta * &f * (1.0 - &f);
        let mut fg = Figure::new();
        let x: Vec<f64> = mu.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = par_f.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/Haldan");
        pdf_name.push_str("/par_f.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        //画一下omega_n 随能量的分布
        let kvec: Array2<f64> = gen_kmesh(&kmesh).unwrap();
        let nk: usize = kvec.len_of(Axis(0));
        let (omega, band) =
            model.berry_curvature_dipole_n(&kvec, &dir_1, &dir_2, &dir_3, og, spin, eta);
        let omega = omega.into_raw_vec();
        let omega = Array1::from(omega);
        let band = band.into_raw_vec();
        let band = Array1::from(band);
        let mut fg = Figure::new();
        let x: Vec<f64> = band.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = omega.to_vec();
        axes.points(
            x.iter(),
            y.iter(),
            &[Color("black"), PointSymbol((".").chars().next().unwrap())],
        );
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/Haldan");
        pdf_name.push_str("/omega_energy.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        //画一下表面态
        let nk = 101;
        let green = surf_Green::from_Model(&model, 0, 1e-3, None).unwrap();
        let E_min = -3.0;
        let E_max = 3.0;
        let E_n = 101;
        let path = [[0.0], [0.5], [1.0]];
        let path = arr2(&path);
        let label = vec!["G", "M", "G"];
        green.show_surf_state("tests/Haldan/surf", &path, &label, nk, E_min, E_max, E_n, 0);

        //-----算一下wilson loop 的结果-----------------------
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let occ = vec![0];
        let wcc = model.wannier_centre(&occ, &array![0.0, 0.0], &dir_1, &dir_2, 101, 101);
        let nocc = occ.len();

        let mut fg = Figure::new();
        let x: Vec<f64> = Array1::<f64>::linspace(0.0, 1.0, 101).to_vec();
        let axes = fg.axes2d();
        for j in -1..2 {
            for i in 0..nocc {
                let a = wcc.row(i).to_owned() + (j as f64) * 2.0 * PI;
                let y: Vec<f64> = a.to_vec();
                axes.points(
                    &x,
                    &y,
                    &[
                        Color("black"),
                        gnuplot::PointSymbol('O'),
                        gnuplot::PointSize(0.2),
                    ],
                );
            }
        }
        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
        let axes = axes.set_y_range(Fix(0.0), Fix(2.0 * PI));
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(0.5, Fix("π")),
            Major(1.0, Fix("2π")),
        ];
        axes.set_x_ticks_custom(
            show_ticks.into_iter(),
            &[],
            &[Font("Times New Roman", 32.0)],
        );
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(PI, Fix("π")),
            Major(2.0 * PI, Fix("2π")),
        ];
        axes.set_y_ticks_custom(
            show_ticks.into_iter(),
            &[],
            &[Font("Times New Roman", 32.0)],
        );
        axes.set_x_label(
            "k_x",
            &[Font("Times New Roman", 32.0), TextOffset(0.0, -0.5)],
        );
        axes.set_y_label(
            "WCC",
            &[
                Font("Times New Roman", 32.0),
                Rotate(90.0),
                TextOffset(-1.0, 0.0),
            ],
        );
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/Haldan/wcc.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();
        //-----------用 berry_flux 算一下
        let C = model
            .berry_flux(
                &occ,
                &array![0.0, 0.0],
                &array![1.0, 0.0],
                &array![0.0, 1.0],
                101,
                101,
            )
            .sum()
            / PI
            / 2.0;
        println!("The Chern number of Haldan model is {}", C);
    }
    #[test]
    fn graphene() {
        let li: Complex<f64> = 1.0 * Complex::i();
        let t1 = 1.0 + 0.0 * li;
        let t2 = 0.0 + 0.0 * li;
        let t3 = 0.0 + 0.0 * li;
        let delta = 0.0;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[3.0_f64.sqrt(), -1.0], [3.0_f64.sqrt(), 1.0]]);
        let orb = arr2(&[[0.0, 0.0], [1.0 / 3.0, 1.0 / 3.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[delta, -delta]), SpinDirection::None);
        model.add_hop(t1, 0, 1, &array![0, 0], SpinDirection::None);
        model.add_hop(t1, 0, 1, &array![-1, 0], SpinDirection::None);
        model.add_hop(t1, 0, 1, &array![0, -1], SpinDirection::None);
        model.add_hop(t2, 0, 0, &array![1, 0], SpinDirection::None);
        model.add_hop(t2, 1, 1, &array![1, 0], SpinDirection::None);
        model.add_hop(t2, 0, 0, &array![0, 1], SpinDirection::None);
        model.add_hop(t2, 1, 1, &array![0, 1], SpinDirection::None);
        model.add_hop(t2, 0, 0, &array![1, -1], SpinDirection::None);
        model.add_hop(t2, 1, 1, &array![1, -1], SpinDirection::None);
        model.add_hop(t3, 0, 1, &array![1, -1], SpinDirection::None);
        model.add_hop(t3, 0, 1, &array![-1, 1], SpinDirection::None);
        model.add_hop(t3, 0, 1, &array![-1, -1], SpinDirection::None);
        let nk: usize = 101;
        let path = [[0.0, 0.0], [2.0 / 3.0, 1.0 / 3.0], [0.5, 0.5], [0.0, 0.0]];
        let path = arr2(&path);
        let (k_vec, k_dist, k_node) = model.k_path(&path, nk).unwrap();
        let (eval, evec) = model.solve_all_parallel(&k_vec);
        let label = vec!["G", "K", "M", "G"];
        model
            .show_band(&path, &label, nk, "tests/graphene")
            .unwrap();

        // 开始计算两个本征态
        let k1 = array![1.0 / 3.0 - 0.002, 2.0 / 3.0];
        let k2 = array![1.0 / 3.0 + 0.001, 2.0 / 3.0];
        let (eval1, evec1) = model.solve_onek(&k1);
        let (eval2, evec2) = model.solve_onek(&k2);
        let evec1 = evec1.reversed_axes();
        let evec2 = evec2.mapv(|x| x.conj());
        println!("{},{}", eval1, eval2);
        println!("{}", evec2.dot(&evec1).mapv(|x| x.norm().round()));

        /////开始计算体系的霍尔电导率//////
        let nk: usize = 11;
        let T: f64 = 0.0;
        let eta: f64 = 0.001;
        let og: f64 = 0.0;
        let mu: f64 = 0.0;
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let spin: usize = 0;
        let kmesh = arr1(&[nk, nk]);
        let (eval, evec) = model.solve_onek(&arr1(&[0.3, 0.5]));
        let conductivity = model.Hall_conductivity(&kmesh, &dir_1, &dir_2, mu, T, spin, eta);
        //println!("{}",conductivity/(2.0*PI));
        //开始计算边缘态, 首先是zigsag态
        let nk: usize = 501;
        let U = arr2(&[[1.0, 1.0], [-1.0, 1.0]]);
        let super_model = model.make_supercell(&U).unwrap();
        let zig_model = super_model.cut_piece(100, 0).unwrap();
        let path = [[0.0, 0.0], [0.0, 0.5], [0.0, 1.0]];
        //let path=[[0.0,0.0],[0.5,0.0],[1.0,0.0]];
        //let path=[[0.0,0.0],[0.5,0.0],[0.5,0.5],[0.0,0.5],[0.0,0.0]];
        let path = arr2(&path);
        let (k_vec, k_dist, k_node) = super_model.k_path(&path, nk).unwrap();
        let (eval, evec) = super_model.solve_all_parallel(&k_vec);
        //let label=vec!["G","X","M","Y","G"];
        let label = vec!["G", "M", "G"];
        zig_model.show_band(&path, &label, nk, "tests/graphene_zig");

        //开始计算石墨烯的态密度
        let nk: usize = 51;
        let kmesh = arr1(&[nk, nk]);
        let E_min = -3.0;
        let E_max = 3.0;
        let E_n = 1000;
        let (E0, dos) = model.dos(&kmesh, E_min, E_max, E_n, 1e-2).unwrap();
        //开始绘制dos
        let mut fg = Figure::new();
        let x: Vec<f64> = E0.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = dos.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/graphene");
        pdf_name.push_str("/dos.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        //开始计算非线性霍尔电导
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let dir_3 = arr1(&[1.0, 0.0]);
        let og = 0.0;
        let mu = Array1::linspace(E_min, E_max, E_n);
        let T = 300.0;
        let sigma: Array1<f64> = model
            .Nonlinear_Hall_conductivity_Extrinsic(
                &kmesh, &dir_1, &dir_2, &dir_3, &mu, T, og, 0, 1e-5,
            )
            .unwrap();

        //开始绘制非线性电导
        let mut fg = Figure::new();
        let x: Vec<f64> = mu.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = sigma.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/graphene");
        pdf_name.push_str("/nonlinear_ex.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();
    }

    #[test]
    fn kane_mele() {
        let li: Complex<f64> = 1.0 * Complex::i();
        let t = -1.0;
        let delta = 0.0;
        let alter = 0.0 + 0.0 * li;
        let soc = 0.06 * t;
        let rashba = 0.0 * t;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[1.0, 0.0], [0.5, 3.0_f64.sqrt() / 2.0]]);
        let orb = arr2(&[[1.0 / 3.0, 1.0 / 3.0], [2.0 / 3.0, 2.0 / 3.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, true, None).unwrap();
        model.set_onsite(&arr1(&[delta, -delta]), SpinDirection::None);
        let R0: Array2<isize> = arr2(&[[0, 0], [-1, 0], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.set_hop(t, 0, 1, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[1, 0], [-1, 1], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.set_hop(soc * li, 0, 0, &R, SpinDirection::z);
        }
        let R0: Array2<isize> = arr2(&[[-1, 0], [1, -1], [0, 1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.set_hop(soc * li, 1, 1, &R, SpinDirection::z);
        }
        //加入rashba项
        let R0: Array2<isize> = arr2(&[[1, 0], [-1, 1], [0, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            let r0 = R.map(|x| *x as f64).dot(&model.lat);
            model.add_hop(rashba * li * r0[[1]], 0, 0, &R, SpinDirection::x);
            model.add_hop(rashba * li * r0[[0]], 0, 0, &R, SpinDirection::y);
        }

        let R0: Array2<isize> = arr2(&[[-1, 0], [1, -1], [0, 1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            let r0 = R.map(|x| *x as f64).dot(&model.lat);
            model.add_hop(-rashba * li * r0[[1]], 1, 1, &R, SpinDirection::x);
            model.add_hop(-rashba * li * r0[[0]], 1, 1, &R, SpinDirection::y);
        }
        let nk: usize = 101;
        let path = [
            [0.0, 0.0],
            [2.0 / 3.0, 1.0 / 3.0],
            [0.5, 0.5],
            [1.0 / 3.0, 2.0 / 3.0],
            [0.0, 0.0],
        ];
        let path = arr2(&path);
        let (k_vec, k_dist, k_node) = model.k_path(&path, nk).unwrap();
        let (eval, evec) = model.solve_all_parallel(&k_vec);
        let label = vec!["G", "K", "M", "K'", "G"];
        model.show_band(&path, &label, nk, "tests/kane");
        //开始计算超胞

        let super_model = model.cut_piece(50, 0).unwrap();
        let path = [[0.0, 0.0], [0.0, 0.5], [0.0, 1.0]];
        let path = arr2(&path);
        let label = vec!["G", "M", "G"];
        super_model
            .show_band(&path, &label, nk, "tests/kane_super")
            .unwrap();
        //开始计算表面态
        let nk = 101;
        let green = surf_Green::from_Model(&model, 0, 1e-3, None).unwrap();
        let E_min = -1.0;
        let E_max = 1.0;
        let E_n = 101;
        let path = [[0.0], [0.5], [1.0]];
        let path = arr2(&path);
        let label = vec!["G", "M", "G"];
        green.show_surf_state("tests/kane", &path, &label, nk, E_min, E_max, E_n, 0);

        //-----算一下wilson loop 结果-----------------------
        let n = 51;
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let occ = vec![0, 1];
        let wcc = model.wannier_centre(&occ, &array![0.0, 0.0], &dir_1, &dir_2, n, n);
        let nocc = occ.len();
        let mut fg = Figure::new();
        let x: Vec<f64> = Array1::<f64>::linspace(0.0, 1.0, n).to_vec();
        let axes = fg.axes2d();
        for j in -1..2 {
            for i in 0..nocc {
                let a = wcc.row(i).to_owned() + (j as f64) * 2.0 * PI;
                let y: Vec<f64> = a.to_vec();
                axes.points(&x, &y, &[Color("black"), gnuplot::PointSymbol('O')]);
            }
        }
        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
        let axes = axes.set_y_range(Fix(0.0), Fix(2.0 * PI));
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(0.5, Fix("π")),
            Major(1.0, Fix("2π")),
        ];
        axes.set_x_ticks_custom(
            show_ticks.into_iter(),
            &[],
            &[Font("Times New Roman", 32.0)],
        );
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(PI, Fix("π")),
            Major(2.0 * PI, Fix("2π")),
        ];
        axes.set_y_ticks_custom(
            show_ticks.into_iter(),
            &[],
            &[Font("Times New Roman", 32.0)],
        );
        axes.set_x_label(
            "k_x",
            &[Font("Times New Roman", 32.0), TextOffset(0.0, -0.5)],
        );
        axes.set_y_label(
            "WCC",
            &[
                Font("Times New Roman", 32.0),
                Rotate(90.0),
                TextOffset(-1.0, 0.0),
            ],
        );
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/kane/wcc.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        /////开始计算体系的霍尔电导率//////
        let nk: usize = 31;
        let T: f64 = 0.0;
        let eta: f64 = 0.001;
        let og: f64 = 0.0;
        let mu: f64 = 0.0;
        //let dir_1=arr1(&[3.0_f64.sqrt()/2.0,-0.5]);
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let spin: usize = 3;
        let kmesh = arr1(&[nk, nk]);
        let start = Instant::now(); // 开始计时
        let conductivity = model
            .Hall_conductivity(&kmesh, &dir_1, &dir_2, mu, T, spin, eta)
            .unwrap();
        let end = Instant::now(); // 结束计时
        let duration = end.duration_since(start); // 计算执行时间
        println!("{}", conductivity * (2.0 * PI));
        println!("function_a took {} seconds", duration.as_secs_f64()); // 输出执行时间
        let nk: usize = 21;
        let kmesh = arr1(&[nk, nk]);
        let start = Instant::now(); // 开始计时
        let conductivity = model
            .Hall_conductivity_adapted(&kmesh, &dir_1, &dir_2, mu, T, spin, eta, 0.01, 0.01)
            .unwrap();
        let end = Instant::now(); // 结束计时
        let duration = end.duration_since(start); // 计算执行时间
        println!("{}", conductivity * (2.0 * PI));
        println!("function_a took {} seconds", duration.as_secs_f64()); // 输出执行时间

        let (E0, dos) = model.dos(&kmesh, E_min, E_max, E_n, 1e-2).unwrap();
        //开始绘制dos
        let mut fg = Figure::new();
        let x: Vec<f64> = E0.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = dos.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/kane");
        pdf_name.push_str("/dos.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();
        //绘制非线性霍尔电导的平面图

        //画一下贝利曲率的分布
        let nk: usize = 31;
        let kmesh = arr1(&[nk, nk]);
        let kvec = gen_kmesh(&kmesh).unwrap();
        //let kvec=kvec-0.5;
        let kvec = kvec * 2.0;
        let kvec = model.lat.dot(&(kvec.reversed_axes()));
        let kvec = kvec.reversed_axes();
        let berry_curv = model.berry_curvature(&kvec, &dir_1, &dir_2, T, 0.0, 1, 1e-3);
        let data = berry_curv.into_shape((nk, nk)).unwrap();
        draw_heatmap(
            &(-data).map(|x| (x + 1.0).log(10.0)),
            "./tests/kane/berry_curvature_distribution.pdf",
        );

        //开始考虑磁场, 加入磁性
        let B = 0.1 + 0.0 * li;
        let tha = 0.0 / 180.0 * PI;

        model.add_hop(B * tha.cos(), 0, 0, &array![0, 0], SpinDirection::x);
        model.add_hop(B * tha.cos(), 1, 1, &array![0, 0], SpinDirection::x);
        model.add_hop(B * tha.sin(), 0, 0, &array![0, 0], SpinDirection::y);
        model.add_hop(B * tha.sin(), 1, 1, &array![0, 0], SpinDirection::y);
        //考虑添加onsite 项破坏空间反演和mirror

        let green = surf_Green::from_Model(&model, 0, 1e-3, None).unwrap();
        let E_min = -1.0;
        let E_max = 1.0;
        let E_n = nk;
        let path = [[0.0], [0.5], [1.0]];
        let path = arr2(&path);
        let label = vec!["G", "M", "G"];
        green.show_surf_state(
            "tests/kane/magnetic",
            &path,
            &label,
            nk,
            E_min,
            E_max,
            E_n,
            0,
        );

        //-----算一下wilson loop 结果-----------------------
        let n = 51;
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let occ = vec![0, 1];
        let wcc = model.wannier_centre(&occ, &array![0.0, 0.0], &dir_1, &dir_2, n, n);
        let nocc = occ.len();
        let mut fg = Figure::new();
        let x: Vec<f64> = Array1::<f64>::linspace(0.0, 1.0, n).to_vec();
        let axes = fg.axes2d();
        for j in -1..2 {
            for i in 0..nocc {
                let a = wcc.row(i).to_owned() + (j as f64) * 2.0 * PI;
                let y: Vec<f64> = a.to_vec();
                axes.points(&x, &y, &[Color("black"), gnuplot::PointSymbol('O')]);
            }
        }
        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
        let axes = axes.set_y_range(Fix(0.0), Fix(2.0 * PI));
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(0.5, Fix("π")),
            Major(1.0, Fix("2π")),
        ];
        axes.set_x_ticks_custom(
            show_ticks.into_iter(),
            &[],
            &[Font("Times New Roman", 32.0)],
        );
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(PI, Fix("π")),
            Major(2.0 * PI, Fix("2π")),
        ];
        axes.set_y_ticks_custom(
            show_ticks.into_iter(),
            &[],
            &[Font("Times New Roman", 32.0)],
        );
        axes.set_x_label(
            "k_x",
            &[Font("Times New Roman", 32.0), TextOffset(0.0, -0.5)],
        );
        axes.set_y_label(
            "WCC",
            &[
                Font("Times New Roman", 32.0),
                Rotate(90.0),
                TextOffset(-1.0, 0.0),
            ],
        );
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/kane/magnetic/wcc.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        //开始计算角态
        let model = model
            .make_supercell(&array![[0.0, -1.0], [1.0, 0.0]])
            .unwrap();
        let num = 19;
        /*
        let model_1=model.cut_piece(num,0).unwrap();
        let new_model=model_1.cut_piece(num,1);
        */
        let new_model = model.cut_dot(num, 6, None).unwrap();
        let mut s = 0;
        let start = Instant::now();
        let (band, evec) = new_model.solve_range_onek(&arr1(&[0.0, 0.0]), (-0.3, 0.3), 1e-5);
        let end = Instant::now(); // 结束计时
        let duration = end.duration_since(start); // 计算执行时间
        println!("solve_band_all took {} seconds", duration.as_secs_f64()); // 输出执行时间
        let nresults = band.len();
        let show_evec = evec.to_owned().map(|x| x.norm_sqr());
        let mut size = Array2::<f64>::zeros((new_model.nsta(), new_model.natom()));
        let norb = new_model.norb();
        for i in 0..nresults {
            let mut s = 0;
            for j in 0..new_model.natom() {
                for k in 0..new_model.atoms[j].norb() {
                    size[[i, j]] += show_evec[[i, s]] + show_evec[[i, s + new_model.norb()]];
                    s += 1;
                }
            }
        }

        let show_str = new_model.atom_position().dot(&model.lat);
        let show_str = show_str.slice(s![.., 0..2]).to_owned();
        let show_size = size.row(new_model.norb()).to_owned();
        create_dir_all("tests/kane/magnetic").expect("can't creat the file");
        write_txt_1(band, "tests/kane/magnetic/band.txt");
        write_txt(size, "tests/kane/magnetic/evec.txt");
        write_txt(show_str, "tests/kane/magnetic/structure.txt");
        //开始绘制角态
    }

    #[test]
    fn Enonlinear() {
        //!arxiv 1706.07702
        //!用来测试非线性外在的霍尔电导
        let li: Complex<f64> = 1.0 * Complex::i();
        let delta = 0.;
        let t1 = 1.0 + 0.0 * li;
        let t2 = 0.2 * t1;
        let t3 = 0.2 * t1;
        let dim_r: usize = 3;
        let norb: usize = 2;
        let lat = arr2(&[
            [1.0, 0.0, 0.0],
            [0.5, 3.0_f64.sqrt() / 2.0, 0.0],
            [0.0, 0.0, 1.0],
        ]);
        let orb = arr2(&[[1.0 / 3.0, 1.0 / 3.0, 0.0], [2.0 / 3.0, 2.0 / 3.0, 0.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[delta, -delta]), SpinDirection::None);
        let R0: Array2<isize> = arr2(&[[0, 0, 0], [-1, 0, 0], [0, -1, 0]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.set_hop(t1, 0, 1, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[1, 0, 1], [-1, 1, 1], [0, -1, 1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.set_hop(t2, 0, 0, &R, SpinDirection::None);
        }
        let R0: Array2<isize> = arr2(&[[1, 0, -1], [-1, 1, -1], [0, -1, -1]]);
        for (i, R) in R0.axis_iter(Axis(0)).enumerate() {
            let R = R.to_owned();
            model.set_hop(t2, 1, 1, &R, SpinDirection::None);
        }
        let R = arr1(&[0, 0, 1]);
        model.set_hop(t3, 0, 0, &R, SpinDirection::None);
        model.set_hop(t3, 1, 1, &R, SpinDirection::None);
        let path = array![
            [0.0, 0.0, 0.0],
            [1.0 / 3.0, 2.0 / 3.0, 0.0],
            [0.5, 0.5, 0.0],
            [0.0, 0.0, 0.0],
            [1.0 / 3.0, 2.0 / 3.0, 0.0],
            [1.0 / 3.0, 2.0 / 3.0, 0.5],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.5],
            [1.0 / 3.0, 2.0 / 3.0, 0.5],
            [0.5, 0.5, 0.5],
            [0.0, 0.0, 0.5]
        ];
        let label = vec!["G", "K", "M", "G", "K", "H", "G", "A", "H", "L", "A"];
        let nk = 101;
        model.show_band(&path, &label, nk, "tests/Enonlinear");

        //开始计算非线性霍尔电导
        let dir_1 = arr1(&[1.0, 0.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0, 0.0]);
        let dir_3 = arr1(&[0.0, 0.0, 1.0]);
        let nk: usize = 21;
        let kmesh = arr1(&[nk, nk, nk]);
        let E_min = -3.0;
        let E_max = 3.0;
        let E_n = 1000;
        let og = 0.0;
        let mu = Array1::linspace(E_min, E_max, E_n);
        let T = 30.0;
        let sigma: Array1<f64> = model
            .Nonlinear_Hall_conductivity_Extrinsic(
                &kmesh, &dir_1, &dir_2, &dir_3, &mu, T, og, 0, 1e-5,
            )
            .unwrap();

        //开始绘制非线性电导
        let mut fg = Figure::new();
        let x: Vec<f64> = mu.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = sigma.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        axes.set_y_range(Fix(-10.0), Fix(10.0));
        axes.set_x_range(Fix(E_min), Fix(E_max));
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/Enonlinear");
        pdf_name.push_str("/nonlinear_ex.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        let sigma: Array1<f64> = model
            .Nonlinear_Hall_conductivity_Intrinsic(&kmesh, &dir_1, &dir_2, &dir_3, &mu, T, 3)
            .unwrap();
        //开始绘制非线性电导
        let mut fg = Figure::new();
        let x: Vec<f64> = mu.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = sigma.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        axes.set_y_range(Fix(-10.0), Fix(10.0));
        axes.set_x_range(Fix(E_min), Fix(E_max));
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/Enonlinear");
        pdf_name.push_str("/nonlinear_in.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();

        let (E0, dos) = model.dos(&kmesh, E_min, E_max, E_n, 1e-2).unwrap();
        //开始绘制dos
        let mut fg = Figure::new();
        let x: Vec<f64> = E0.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = dos.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/Enonlinear");
        pdf_name.push_str("/dos.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();
    }
    #[test]
    fn kagome() {
        let li: Complex<f64> = 1.0 * Complex::i();
        let t1 = 1.0 + 0.0 * li;
        let t2 = 0.1 + 0.0 * li;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[3.0_f64.sqrt(), -1.0], [3.0_f64.sqrt(), 1.0]]);
        let orb = arr2(&[[0.0, 0.0], [1.0 / 3.0, 0.0], [0.0, 1.0 / 3.0]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        //最近邻hopping
        model.add_hop(t1, 0, 1, &array![0, 0], SpinDirection::None);
        model.add_hop(t1, 2, 0, &array![0, 0], SpinDirection::None);
        model.add_hop(t1, 1, 2, &array![0, 0], SpinDirection::None);
        model.add_hop(t1, 0, 2, &array![0, -1], SpinDirection::None);
        model.add_hop(t1, 0, 1, &array![-1, 0], SpinDirection::None);
        model.add_hop(t1, 2, 1, &array![-1, 1], SpinDirection::None);
        let nk: usize = 101;
        let path = [[0.0, 0.0], [2.0 / 3.0, 1.0 / 3.0], [0.5, 0.], [0.0, 0.0]];
        let path = arr2(&path);
        let label = vec!["G", "K", "M", "G"];
        model.show_band(&path, &label, nk, "tests/kagome/");
        //start to draw the band structure
        //Starting to calculate the edge state, first is the zigzag state
        let nk: usize = 101;
        let U = arr2(&[[1.0, 1.0], [-1.0, 1.0]]);
        let super_model = model.make_supercell(&U).unwrap();
        let zig_model = super_model.cut_piece(30, 0).unwrap();
        let path = [[0.0, 0.0], [0.0, 0.5], [0.0, 1.0]];
        let path = arr2(&path);
        let (k_vec, k_dist, k_node) = super_model.k_path(&path, nk).unwrap();
        let (eval, evec) = super_model.solve_all_parallel(&k_vec);
        let label = vec!["G", "M", "G"];
        zig_model.show_band(&path, &label, nk, "tests/kagome_zig/");

        let green = surf_Green::from_Model(&super_model, 0, 1e-3, None).unwrap();
        let E_min = -2.0;
        let E_max = 4.0;
        let E_n = nk;
        let path = [[0.0], [0.5], [1.0]];
        let path = arr2(&path);
        let label = vec!["G", "M", "G"];
        green.show_surf_state("tests/kagome_zig", &path, &label, nk, E_min, E_max, E_n, 0);

        //Starting to calculate the DOS of kagome
        let nk: usize = 51;
        let kmesh = arr1(&[nk, nk]);
        let E_min = -3.0;
        let E_max = 3.0;
        let E_n = 1000;
        let (E0, dos) = model.dos(&kmesh, E_min, E_max, E_n, 1e-2).unwrap();
        //start to show DOS
        let mut fg = Figure::new();
        let x: Vec<f64> = E0.to_vec();
        let axes = fg.axes2d();
        let y: Vec<f64> = dos.to_vec();
        axes.lines(&x, &y, &[Color("black")]);
        let mut show_ticks = Vec::<String>::new();
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/kagome/");
        pdf_name.push_str("dos.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();
    }

    #[test]
    fn SSH() {
        let li: Complex<f64> = 1.0 * Complex::i();
        let t1 = 1.0 + 0.0 * li;
        let t2 = 0.5 + 0.0 * li;
        let Delta = 0.0;
        let dim_r: usize = 1;
        let norb: usize = 2;
        let lat = arr2(&[[1.0]]);
        let orb = arr2(&[[0.3], [0.5]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.add_hop(t1, 0, 1, &array![0], SpinDirection::None);
        model.add_hop(t2, 0, 1, &array![-1], SpinDirection::None);
        model.add_onsite(&array![Delta, -Delta], 0);

        let nk: usize = 101;
        let path = [[0.0], [0.5], [1.0]];
        let path = arr2(&path);
        let label = vec!["G", "M", "G"];
        model.show_band(&path, &label, nk, "tests/SSH/");
        let mut super_model = model.cut_piece(5, 0).unwrap();

        let (band, evec) = super_model.solve_onek(&array![0.0]);
        println!("{}", band);
    }
    #[test]
    fn BBH_model() {
        let li: Complex<f64> = 1.0 * Complex::i();
        let t1 = 0.1 + 0.0 * li;
        let t2 = 1.0 + 0.0 * li;
        let i0 = -1.0;
        let dim_r: usize = 2;
        let norb: usize = 2;
        let lat = arr2(&[[1.0, 0.0], [0.0, 1.0]]);
        let orb = arr2(&[[0.0, 0.0], [0.5, 0.0], [0.5, 0.5], [0.0, 0.5]]);
        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.add_hop(t1, 0, 1, &array![0, 0], SpinDirection::None);
        model.add_hop(t1, 1, 2, &array![0, 0], SpinDirection::None);
        model.add_hop(t1, 2, 3, &array![0, 0], SpinDirection::None);
        model.add_hop(i0 * t1, 3, 0, &array![0, 0], SpinDirection::None);
        model.add_hop(t2, 0, 1, &array![-1, 0], SpinDirection::None);
        model.add_hop(i0 * t2, 0, 3, &array![0, -1], SpinDirection::None);
        model.add_hop(t2, 2, 3, &array![1, 0], SpinDirection::None);
        model.add_hop(t2, 2, 1, &array![0, 1], SpinDirection::None);
        let nk: usize = 101;
        let path = [[0.0, 0.0], [0.5, 0.0], [0.5, 0.5], [0.0, 0.0]];
        let path = arr2(&path);
        let label = vec!["G", "X", "M", "G"];
        model.show_band(&path, &label, nk, "tests/BBH/").unwrap();
        model.output_hr("tests/BBH/", "wannier90");

        //算一下wilson loop
        let n = 51;
        let dir_1 = arr1(&[1.0, 0.0]);
        let dir_2 = arr1(&[0.0, 1.0]);
        let occ = vec![0, 1];
        let wcc = model.wannier_centre(&occ, &array![0.0, 0.0], &dir_1, &dir_2, n, n);
        let nocc = occ.len();
        let mut fg = Figure::new();
        let x: Vec<f64> = Array1::<f64>::linspace(0.0, 1.0, n).to_vec();
        let axes = fg.axes2d();
        for j in -1..2 {
            for i in 0..nocc {
                let a = wcc.row(i).to_owned() + (j as f64) * 2.0 * PI;
                let y: Vec<f64> = a.to_vec();
                axes.points(&x, &y, &[Color("black"), gnuplot::PointSymbol('O')]);
            }
        }
        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
        let axes = axes.set_y_range(Fix(0.0), Fix(2.0 * PI));
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(0.5, Fix("π")),
            Major(1.0, Fix("2π")),
        ];
        axes.set_x_ticks_custom(show_ticks.into_iter(), &[], &[]);
        let show_ticks = vec![
            Major(0.0, Fix("0")),
            Major(PI, Fix("π")),
            Major(2.0 * PI, Fix("2π")),
        ];
        axes.set_y_ticks_custom(show_ticks.into_iter(), &[], &[]);
        let mut pdf_name = String::new();
        pdf_name.push_str("tests/BBH/wcc.pdf");
        fg.set_terminal("pdfcairo", &pdf_name);
        fg.show();
        //算一下边界态
        let green = surf_Green::from_Model(&model, 0, 1e-3, None).unwrap();
        let E_min = -2.0;
        let E_max = 2.0;
        let E_n = nk;
        let path = [[0.0], [0.5], [1.0]];
        let path = arr2(&path);
        let label = vec!["G", "X", "G"];
        green.show_surf_state("tests/BBH", &path, &label, nk, E_min, E_max, E_n, 0);

        //算一下corner state
        let num = 10;
        let model_1 = model.cut_piece(num, 0).unwrap();
        let new_model = model_1.cut_piece(2 * num, 1).unwrap();
        let mut s = 0;
        let start = Instant::now();
        let (band, evec) = new_model.solve_onek(&arr1(&[0.0, 0.0]));
        println!(
            "band shape is {:?}, evec shape is {:?}",
            band.shape(),
            evec.shape()
        );
        let end = Instant::now(); // 结束计时
        let duration = end.duration_since(start); // 计算执行时间
        println!("solve_band_all took {} seconds", duration.as_secs_f64()); // 输出执行时间
        let nresults = band.len();
        let show_evec = evec.to_owned().map(|x| x.norm_sqr());
        let norb = new_model.norb();
        let size = show_evec;
        let show_str = new_model.atom_position().dot(&model.lat);
        create_dir_all("tests/BBH/corner").expect("can't creat the file");
        write_txt_1(band, "tests/BBH/corner/band.txt");
        write_txt(size, "tests/BBH/corner/evec.txt");
        write_txt(show_str, "tests/BBH/corner/structure.txt");
    }

    #[test]
    fn graphene_magnetic_field() {
        use crate::{MagneticField, Model, SpinDirection};
        use ndarray::{Axis, arr1, arr2};
        use num_complex::Complex;
        // 如果你在其他地方定义了画图函数，请确保 use 进来，例如：
        // use crate::draw_heatmap;

        // 1. 设置模型基本参数
        let t = Complex::new(-1.0, 0.0);
        let delta = 0.0;
        let dim_r: usize = 2;
        let norb: usize = 2;

        // 石墨烯晶格：a1 = (1, 0), a2 = (1/2, √3/2)
        let lat = arr2(&[[1.0, 0.0], [0.5, 3.0_f64.sqrt() / 2.0]]);
        // 轨道的相对分数坐标 (Fractional Coordinates)
        let orb = arr2(&[[1.0 / 3.0, 1.0 / 3.0], [2.0 / 3.0, 2.0 / 3.0]]);

        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[-delta, delta]), SpinDirection::None);

        // 添加最近邻跃迁
        let r0: ndarray::Array2<isize> = arr2(&[[0, 0], [-1, 0], [0, -1]]);
        for r in r0.axis_iter(Axis(0)) {
            model.add_hop(t, 0, 1, &r.to_owned(), SpinDirection::None);
        }

        // 2. 施加磁场
        // 重要：二维中，面外磁场垂直于 xy 平面，对应的索引必须为 z 轴 (即 mag_dir = 2)
        // 扩胞 [3, 3] 表示 a1 和 a2 两个方向各扩胞 3 倍
        // 磁通 8 表示整体超胞含有 8 个磁通量子 (每个原胞 8/9 个磁通)
        let magnetic_model = model.add_magnetic_field(2, [9, 9], 40).unwrap();

        // 3. 高对称路径 (注意：六角晶格的 K 点在分数倒格矢坐标下是 1/3, 1/3)
        let path = arr2(&[[0.0, 0.0], [0.5, 0.0], [1.0 / 3.0, 1.0 / 3.0], [0.0, 0.0]]);
        let label = vec!["Γ", "M", "K", "Γ"];
        let nk = 1001;

        // 4. 绘制折叠态下的超胞能带 (Hofstadter 蝴蝶状能带切片)
        magnetic_model
            .show_band(&path, &label, nk, "tests/graphene_magnetic")
            .unwrap();

        // 5. 展开能带 (Unfold) 回到原胞 Brillouin 区
        let u_matrix = arr2(&[[9.0, 0.0], [0.0, 9.0]]);

        // 生成对应的谱函数 (Spectral Weight)
        let a_spectral = magnetic_model
            .unfold(&u_matrix, &path, nk, -3.0, 3.0, nk, 1e-3, 1e-5)
            .unwrap();

        // 将谱函数输出为热力图
        // (假定 draw_heatmap 接收二维热力矩阵以及保存路径)
        draw_heatmap(
            &a_spectral.reversed_axes(),
            "./tests/graphene_magnetic/unfold_band.pdf",
        );
    }

    #[test]
    fn test_hofstadter_butterfly_gnuplot() {
        use gnuplot::AutoOption::Fix;
        use gnuplot::{AxesCommon, Color, Figure, PointSize, PointSymbol};

        // 1. 初始化一个标准的 2D 正方晶格
        let t = Complex::new(-1.0, 0.0);
        let dim_r = 2;
        // 正方晶格基矢 a1=(1,0), a2=(0,1)
        let lat = arr2(&[[1.0, 0.0], [0.0, 1.0]]);
        // 单轨道坐标位于原点
        let orb = arr2(&[[0.0, 0.0]]);

        let mut model = Model::tb_model(dim_r, lat, orb, false, None).unwrap();
        model.set_onsite(&arr1(&[0.0]), SpinDirection::None);

        // 加上最近邻跃迁 (上下左右4个方向)
        let r0: ndarray::Array2<isize> = arr2(&[[1, 0], [-1, 0], [0, 1], [0, -1]]);
        for r in r0.axis_iter(Axis(0)) {
            model.add_hop(t, 0, 0, &r.to_owned(), SpinDirection::None);
        }

        // 2. 准备扫描磁通并记录坐标
        let q = 81; // 分母 q 设定为 49 形成的蝴蝶分形已经足够好看

        // 我们用两个动态数组分别记录散点图的 X 轴 (磁通) 和 Y 轴 (能量)
        let mut x_data = Vec::new();
        let mut y_data = Vec::new();

        println!("开始计算 Hofstadter 蝴蝶能谱，进度: ");

        for p in 0..=q {
            // 在 y 轴方向扩胞 q 倍，即 [1, q]，形成一个一维超胞
            let mag_model = model.add_magnetic_field(2, [1, q], p as isize).unwrap();
            let norb = mag_model.norb();

            // 提取 Gamma 点 k = (0,0) 的哈密顿量矩阵
            let mut h_k0 = Array2::<Complex<f64>>::zeros((norb, norb));
            for iR in 0..mag_model.hamR.nrows() {
                for i in 0..norb {
                    for j in 0..norb {
                        h_k0[[i, j]] += mag_model.ham[[iR, i, j]];
                    }
                }
            }

            // --- 求解本征值 ---
            let evals = h_k0.eigvalsh(UPLO::Upper).expect("矩阵对角化失败");

            // 收集坐标点
            let flux_ratio = (p as f64) / (q as f64);
            for &e in evals.iter() {
                x_data.push(flux_ratio);
                y_data.push(e);
            }

            if p % 10 == 0 {
                println!("已完成 {}/{}", p, q);
            }
        }

        println!(
            "计算完成，共有 {} 个能级点，正在使用 gnuplot 绘图...",
            x_data.len()
        );

        // 3. 使用 Gnuplot 直接输出图像
        // 确保目录存在
        create_dir_all("tests").expect("无法创建 tests 文件夹");

        let mut fg = Figure::new();
        let axes = fg.axes2d();

        axes.set_title("Hofstadter's Butterfly", &[]);
        axes.set_x_label("Magnetic Flux (\\Phi / \\Phi_0)", &[]);
        axes.set_y_label("Energy (E/t)", &[]);

        // 固定 x 和 y 的坐标范围
        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
        let axes = axes.set_y_range(Fix(-10.0), Fix(10.0));

        // 使用 .points 绘制散点图
        // PointSymbol('.') 表示画极小的像素点，最适合用来展示密集的分形结构
        // Color("navy") 使用深蓝色
        axes.points(
            &x_data,
            &y_data,
            &[Color("navy"), PointSymbol('.'), PointSize(0.6)],
        );

        // 将图像渲染为 PDF
        fg.set_terminal("pdfcairo", "tests/hofstadter_butterfly.pdf");
        fg.show().expect("Gnuplot 画图失败");

        println!("完美！图像已保存至 tests/hofstadter_butterfly.pdf");
    }
}