Rustb/
surfgreen.rs

1//! Surface Green's function calculations for semi-infinite systems.
2//!
3//! This module implements algorithms for computing the surface Green's function
4//! $G^s(\omega, \mathbf{k}_\parallel)$ of semi-infinite tight-binding systems.
5//! The surface Green's function is defined as:
6//! $$
7//! G^s(\omega, \mathbf{k}_\parallel) = [\omega - H_{00}(\mathbf{k}_\parallel) - \Sigma(\omega, \mathbf{k}_\parallel)]^{-1}
8//! $$
9//! where $\Sigma$ is the self-energy due to the semi-infinite bulk.
10use crate::error::{TbError, Result};
11use crate::{Model};
12use crate::kpoints::gen_kmesh;
13use crate::basis::Dimension;
14use gnuplot::Major;
15use gnuplot::{Auto, AutoOption::Fix, AxesCommon, Custom, Figure, Font, HOT, RAINBOW};
16use ndarray::concatenate;
17use ndarray::linalg::kron;
18use ndarray::prelude::*;
19use ndarray::*;
20use ndarray_linalg::conjugate;
21use ndarray_linalg::*;
22use num_complex::Complex;
23use rayon::prelude::*;
24use std::f64::consts::PI;
25use std::fs::File;
26use std::io::Write;
27use std::ops::AddAssign;
28use std::ops::MulAssign;
29use std::time::Instant;
30
31///计算表面格林函数的时候使用的基本单位
32#[derive(Clone, Debug)]
33pub struct surf_Green {
34    /// - The real space dimension of the model.
35    pub dim_r: usize,
36    /// - The number of orbitals in the model.
37    pub norb: usize,
38    /// - The number of states in the model. If spin is enabled, nsta=norb$\times$2
39    pub nsta: usize,
40    /// - The number of atoms in the model. The atom and atom_list at the back are used to store the positions of the atoms, and the number of orbitals corresponding to each atom.
41    pub natom: usize,
42    /// - Whether the model has spin enabled. If enabled, spin=true
43    pub spin: bool,
44    /// - The lattice vector of the model, a dim_r$\times$dim_r matrix, the axis0 direction stores a 1$\times$dim_r lattice vector.
45    pub lat: Array2<f64>,
46    /// - The position of the orbitals in the model. We use fractional coordinates uniformly.
47    pub orb: Array2<f64>,
48    /// - The position of the atoms in the model, also in fractional coordinates.
49    pub atom: Array2<f64>,
50    /// - The number of orbitals in the atoms, in the same order as the atom positions.
51    pub atom_list: Vec<usize>,
52    /// - The bulk Hamiltonian of the model, $\bra{m0}\hat H\ket{nR}$, a three-dimensional complex tensor of size n_R$\times$nsta$\times$ nsta, where the first nsta*nsta matrix corresponds to hopping within the unit cell, i.e. <m0|H|n0>, and the subsequent matrices correspond to hopping within hamR.
53    pub eta: f64,
54    pub ham_bulk: Array3<Complex<f64>>,
55    /// - The distance between the unit cell hoppings, i.e. R in $\bra{m0}\hat H\ket{nR}$.
56    pub ham_bulkR: Array2<isize>,
57    /// - The bulk Hamiltonian of the model, $\bra{m0}\hat H\ket{nR}$, a three-dimensional complex tensor of size n_R$\times$nsta$\times$ nsta, where the first nsta*nsta matrix corresponds to hopping within the unit cell, i.e. <m0|H|n0>, and the subsequent matrices correspond to hopping within hamR.
58    pub ham_hop: Array3<Complex<f64>>,
59    pub ham_hopR: Array2<isize>,
60}
61
62
63#[inline(always)]
64fn remove_row<T: Copy>(array: Array2<T>, row_to_remove: usize) -> Array2<T> {
65    let indices: Vec<_> = (0..array.nrows()).filter(|&r| r != row_to_remove).collect();
66    array.select(Axis(0), &indices)
67}
68#[inline(always)]
69fn remove_col<T: Copy>(array: Array2<T>, col_to_remove: usize) -> Array2<T> {
70    let indices: Vec<_> = (0..array.ncols()).filter(|&r| r != col_to_remove).collect();
71    array.select(Axis(1), &indices)
72}
73
74
75impl surf_Green {
76    ///从 Model 中构建一个 surf_green 的结构体
77    ///
78    ///dir表示要看哪方向的表面态
79    ///
80    ///eta表示小虚数得取值
81    ///
82    ///对于非晶格矢量得方向, 需要用 model.make_supercell 先扩胞
83    pub fn from_Model(model: &Model, dir: usize, eta: f64, Np: Option<usize>) -> Result<surf_Green> {
84        if dir >= model.dim_r as usize {
85            return Err(TbError::InvalidDirection {
86                index: dir,
87                dim: model.dim_r as usize,
88            });
89        }
90        let mut R_max: usize = 0;
91        for R0 in model.hamR.rows() {
92            if R_max < R0[[dir]].abs() as usize {
93                R_max = R0[[dir]].abs() as usize;
94            }
95        }
96        let R_max = match Np {
97            Some(np) => {
98                if R_max > np {
99                    np
100                } else {
101                    R_max
102                }
103            }
104            None => R_max,
105        };
106
107        let mut U = Array2::<f64>::eye(model.dim_r as usize);
108        U[[dir, dir]] = R_max as f64;
109        let model = model.make_supercell(&U)?;
110        let mut ham0 = Array3::<Complex<f64>>::zeros((0, model.nsta(), model.nsta()));
111        let mut hamR0 = Array2::<isize>::zeros((0, model.dim_r as usize));
112        let mut hamR = Array3::<Complex<f64>>::zeros((0, model.nsta(), model.nsta()));
113        let mut hamRR = Array2::<isize>::zeros((0, model.dim_r as usize));
114        let use_hamR = model.hamR.rows();
115        let use_ham = model.ham.outer_iter();
116        for (ham, R) in use_ham.zip(use_hamR) {
117            let ham = ham.clone();
118            let R = R.clone();
119            if R[[dir]] == 0 {
120                ham0.push(Axis(0), ham.view());
121                hamR0.push_row(R.view());
122            } else if R[[dir]] > 0 {
123                hamR.push(Axis(0), ham.view());
124                hamRR.push_row(R.view());
125            }
126        }
127        let new_lat = remove_row(model.lat.clone(), dir);
128        let new_lat = remove_col(new_lat.clone(), dir);
129        let new_orb = remove_col(model.orb.clone(), dir);
130        let new_atom = remove_col(model.atom_position(), dir);
131        let new_hamR0 = remove_col(hamR0, dir);
132        let new_hamRR = remove_col(hamRR, dir);
133        let green: surf_Green = surf_Green {
134            dim_r: model.dim_r as usize - 1,
135            norb: model.norb(),
136            nsta: model.nsta(),
137            natom: model.natom(),
138            spin: model.spin,
139            lat: new_lat,
140            orb: new_orb,
141            atom: new_atom,
142            atom_list: model.atom_list(),
143            ham_bulk: ham0,
144            ham_bulkR: new_hamR0,
145            ham_hop: hamR,
146            ham_hopR: new_hamRR,
147            eta,
148        };
149        Ok(green)
150    }
151    pub fn k_path(&self, path: &Array2<f64>, nk: usize) -> Result<(Array2<f64>, Array1<f64>, Array1<f64>)> {
152        //!根据高对称点来生成高对称路径, 画能带图
153        if self.dim_r == 0 {
154            return Err(TbError::ZeroDimKPathError);
155        }
156        let n_node: usize = path.len_of(Axis(0));
157        if self.dim_r != path.len_of(Axis(1)) {
158            return Err(TbError::PathLengthMismatch {
159                expected: self.dim_r,
160                actual: path.len_of(Axis(1)),
161            });
162        }
163        let k_metric = (self.lat.dot(&self.lat.t())).inv().unwrap();
164        let mut k_node = Array1::<f64>::zeros(n_node);
165        for n in 1..n_node {
166            let dk = &path.row(n) - &path.row(n - 1);
167            let a = k_metric.dot(&dk);
168            let dklen = dk.dot(&a).sqrt();
169            k_node[[n]] = k_node[[n - 1]] + dklen;
170        }
171        let mut node_index: Vec<usize> = vec![0];
172        for n in 1..n_node - 1 {
173            let frac = k_node[[n]] / k_node[[n_node - 1]];
174            let a = (frac * ((nk - 1) as f64).round()) as usize;
175            node_index.push(a)
176        }
177        node_index.push(nk - 1);
178        let mut k_dist = Array1::<f64>::zeros(nk);
179        let mut k_vec = Array2::<f64>::zeros((nk, self.dim_r));
180        //k_vec.slice_mut(s![0,..]).assign(&path.slice(s![0,..]));
181        k_vec.row_mut(0).assign(&path.row(0));
182        for n in 1..n_node {
183            let n_i = node_index[n - 1];
184            let n_f = node_index[n];
185            let kd_i = k_node[[n - 1]];
186            let kd_f = k_node[[n]];
187            let k_i = path.row(n - 1);
188            let k_f = path.row(n);
189            for j in n_i..n_f + 1 {
190                let frac: f64 = ((j - n_i) as f64) / ((n_f - n_i) as f64);
191                k_dist[[j]] = kd_i + frac * (kd_f - kd_i);
192                k_vec
193                    .row_mut(j)
194                    .assign(&((1.0 - frac) * &k_i + frac * &k_f));
195            }
196        }
197        Ok((k_vec, k_dist, k_node))
198    }
199
200    #[inline(always)]
201    pub fn gen_ham_onek<S: Data<Elem = f64>>(
202        &self,
203        kvec: &ArrayBase<S, Ix1>,
204    ) -> (Array2<Complex<f64>>, Array2<Complex<f64>>) {
205        let mut ham0k = Array2::<Complex<f64>>::zeros((self.nsta, self.nsta));
206        let mut hamRk = Array2::<Complex<f64>>::zeros((self.nsta, self.nsta));
207        if kvec.len() != self.dim_r {
208            panic!("Wrong, the k-vector's length must equal to the dimension of model.")
209        }
210        let nR: usize = self.ham_bulkR.len_of(Axis(0));
211        let nRR: usize = self.ham_hopR.len_of(Axis(0));
212        let U0: Array1<f64> = self.orb.dot(kvec);
213        let U0: Array1<Complex<f64>> = U0.map(|x| Complex::<f64>::new(0.0, 2.0 * PI * *x));
214        let U0: Array1<Complex<f64>> = U0.mapv(Complex::exp); //关于轨道的 guage
215        let U0 = if self.spin {
216            let U0 = concatenate![Axis(0), U0, U0];
217            U0
218        } else {
219            U0
220        };
221        let U: Array2<Complex<f64>> = Array2::from_diag(&U0);
222        let U_conj = Array2::from_diag(&U0.map(|x| x.conj()));
223        //对体系作傅里叶变换
224        let U0 = (self.ham_bulkR.map(|x| *x as f64))
225            .dot(kvec)
226            .map(|x| Complex::<f64>::new(0.0, *x * 2.0 * PI).exp());
227        //对 ham_hop 作傅里叶变换
228        let UR = (self.ham_hopR.map(|x| *x as f64))
229            .dot(kvec)
230            .map(|x| Complex::<f64>::new(0.0, *x * 2.0 * PI).exp());
231        let ham0k = self
232            .ham_bulk
233            .outer_iter()
234            .zip(U0.iter())
235            .fold(Array2::zeros((self.nsta, self.nsta)), |acc, (ham, u)| {
236                acc + &ham * *u
237            });
238        let hamRk = self
239            .ham_hop
240            .outer_iter()
241            .zip(UR.iter())
242            .fold(Array2::zeros((self.nsta, self.nsta)), |acc, (ham, u)| {
243                acc + &ham * *u
244            });
245        //先对 ham_bulk 中的 [0,0] 提取出来
246        //let ham0 = self.ham_bulk.slice(s![0, .., ..]);
247        //let U0 = U0.slice(s![1..nR]);
248        //let U0 = U0.into_shape((nR - 1, 1, 1)).unwrap();
249        //let U0 = U0.broadcast((nR - 1, self.nsta, self.nsta)).unwrap();
250        //let ham0k = (&self.ham_bulk.slice(s![1..nR, .., ..]) * &U0).sum_axis(Axis(0));
251        //let UR = UR.into_shape((nRR, 1, 1)).unwrap();
252        //let UR = UR.broadcast((nRR, self.nsta, self.nsta)).unwrap();
253        //let hamRk = (&self.ham_hop * &UR).sum_axis(Axis(0));
254        //let ham0k: Array2<Complex<f64>> = &ham0
255        //    + &conjugate::<Complex<f64>, OwnedRepr<Complex<f64>>, OwnedRepr<Complex<f64>>>(&ham0k)
256        //    + &ham0k;
257        //作规范变换
258        let ham0k = ham0k.dot(&U);
259        let ham0k = U_conj.dot(&ham0k);
260        let hamRk = hamRk.dot(&U);
261        let hamRk = U_conj.dot(&hamRk);
262        (ham0k, hamRk)
263    }
264    #[inline(always)]
265    pub fn surf_green_one<S: Data<Elem = f64>>(
266        &self,
267        kvec: &ArrayBase<S, Ix1>,
268        Energy: f64,
269    ) -> (f64, f64, f64) {
270        let (hamk, hamRk) = self.gen_ham_onek(kvec);
271        let hamRk_conj: Array2<Complex<f64>> =
272            conjugate::<Complex<f64>, OwnedRepr<Complex<f64>>, OwnedRepr<Complex<f64>>>(&hamRk);
273        let I0 = Array2::<Complex<f64>>::eye(self.nsta);
274        let accurate: f64 = 1e-8;
275        let epsilon = Complex::new(Energy, self.eta) * &I0;
276        let mut epi = hamk.clone();
277        let mut eps = hamk.clone();
278        let mut eps_t = hamk.clone();
279        let mut ap = hamRk.clone();
280        let mut bt = hamRk_conj.clone();
281
282        for _ in 0..10 {
283            let g0 = (&epsilon - &epi).inv().unwrap();
284            let mat_1 = &ap.dot(&g0);
285            let mat_2 = &bt.dot(&g0);
286            let g0 = &mat_1.dot(&bt);
287            epi = epi + g0;
288            eps = eps + g0;
289            let g0 = &mat_2.dot(&ap);
290            epi = epi + g0;
291            eps_t = eps_t + g0;
292            ap = mat_1.dot(&ap);
293            bt = mat_2.dot(&bt);
294            if ap.sum().norm() < accurate {
295                break;
296            }
297        }
298        let g_LL = (&epsilon - eps).inv().unwrap();
299        let g_RR = (&epsilon - eps_t).inv().unwrap();
300        let g_B = (&epsilon - epi).inv().unwrap();
301        let N_R: f64 = -1.0 / (PI) * g_RR.into_diag().sum().im;
302        let N_L: f64 = -1.0 / (PI) * g_LL.into_diag().sum().im;
303        let N_B: f64 = -1.0 / (PI) * g_B.into_diag().sum().im;
304        (N_R, N_L, N_B)
305    }
306
307    #[inline(always)]
308    pub fn surf_green_onek<S: Data<Elem = f64>>(
309        &self,
310        kvec: &ArrayBase<S, Ix1>,
311        Energy: &Array1<f64>,
312    ) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
313        let (hamk, hamRk) = self.gen_ham_onek(kvec);
314        let hamRk_conj: Array2<Complex<f64>> =
315            conjugate::<Complex<f64>, OwnedRepr<Complex<f64>>, OwnedRepr<Complex<f64>>>(&hamRk);
316        let I0 = Array2::<Complex<f64>>::eye(self.nsta);
317        let accurate: f64 = 1e-6;
318        let ((N_R, N_L), N_B): ((Vec<_>, Vec<_>), Vec<_>) = Energy
319            .map(|e| {
320                let epsilon = Complex::new(*e, self.eta) * &I0;
321                let mut epi = hamk.clone();
322                let mut eps = hamk.clone();
323                let mut eps_t = hamk.clone();
324                let mut ap = hamRk.clone();
325                let mut bt = hamRk_conj.clone();
326                for _ in 0..10 {
327                    let g0 = (&epsilon - &epi).inv().unwrap();
328                    let mat_1 = &ap.dot(&g0);
329                    let mat_2 = &bt.dot(&g0);
330                    let g0 = &mat_1.dot(&bt);
331                    epi += g0;
332                    eps += g0;
333                    let g0 = &mat_2.dot(&ap);
334                    epi += g0;
335                    eps_t += g0;
336                    ap = mat_1.dot(&ap);
337                    bt = mat_2.dot(&bt);
338                    if ap.map(|x| x.norm()).sum() < accurate {
339                        break;
340                    }
341                }
342                let g_LL = (&epsilon - eps).inv().unwrap();
343                let g_RR = (&epsilon - eps_t).inv().unwrap();
344                let g_B = (&epsilon - epi).inv().unwrap();
345                //求trace
346                let N_R: f64 = -1.0 / (PI) * g_RR.into_diag().sum().im;
347                let N_L: f64 = -1.0 / (PI) * g_LL.into_diag().sum().im;
348                let N_B: f64 = -1.0 / (PI) * g_B.into_diag().sum().im;
349                ((N_R, N_L), N_B)
350            })
351            .into_iter()
352            .unzip();
353        let N_R = Array1::from_vec(N_R);
354        let N_L = Array1::from_vec(N_L);
355        let N_B = Array1::from_vec(N_B);
356        (N_R, N_L, N_B)
357    }
358
359    pub fn surf_green_path(
360        &self,
361        kvec: &Array2<f64>,
362        E_min: f64,
363        E_max: f64,
364        E_n: usize,
365        spin: usize,
366    ) -> (Array2<f64>, Array2<f64>, Array2<f64>) {
367        let Energy = Array1::<f64>::linspace(E_min, E_max, E_n);
368        let nk = kvec.nrows();
369        let mut N_R = Array2::<f64>::zeros((nk, E_n));
370        let mut N_L = Array2::<f64>::zeros((nk, E_n));
371        let mut N_B = Array2::<f64>::zeros((nk, E_n));
372        Zip::from(N_R.outer_iter_mut())
373            .and(N_L.outer_iter_mut())
374            .and(N_B.outer_iter_mut())
375            .and(kvec.outer_iter())
376            .par_for_each(|mut nr, mut nl, mut nb, k| {
377                let (NR, NL, NB) = self.surf_green_onek(&k, &Energy);
378                nr.assign(&NR);
379                nl.assign(&NL);
380                nb.assign(&NB);
381            });
382        (N_L, N_R, N_B)
383    }
384
385    pub fn show_arc_state(&self, name: &str, kmesh: &Array1<usize>, energy: f64, spin: usize) {
386        use std::fs::create_dir_all;
387        use std::io::{BufWriter, Write};
388        create_dir_all(name).expect("can't creat the file");
389        assert_eq!(
390            kmesh.len(),
391            2,
392            "show_arc_state can only calculated the three dimension system, so the kmesh need to be [m,n], but you give {}",
393            kmesh
394        );
395        let kvec = gen_kmesh(kmesh).expect("Failed to generate k-mesh");
396        let nk = kvec.nrows();
397        let mut N_R = Array1::<f64>::zeros(nk);
398        let mut N_L = Array1::<f64>::zeros(nk);
399        let mut N_B = Array1::<f64>::zeros(nk);
400        Zip::from(N_R.view_mut())
401            .and(N_L.view_mut())
402            .and(N_B.view_mut())
403            .and(kvec.outer_iter())
404            .par_for_each(|mut nr, mut nl, mut nb, k| {
405                let (NR, NL, NB) = self.surf_green_one(&k, energy);
406                *nr = NR;
407                *nl = NL;
408                *nb = NB;
409            });
410        let K = 2.0 * PI * self.lat.inv().unwrap().reversed_axes();
411        let kvec_real = kvec.dot(&K);
412        let mut file_name = String::new();
413        file_name.push_str(&name);
414        file_name.push_str("/arc.dat");
415        let mut file = File::create(file_name).expect("Uable to create arc.dat");
416        writeln!(file, r"# nk1, nk2, N_L, N_R, N_B");
417        let mut writer = BufWriter::new(file);
418        let mut s = String::new();
419        for i in 0..nk {
420            let aa = format!("{:.6}", kvec_real[[i, 0]]);
421            s.push_str(&aa);
422            let bb: String = format!("{:.6}", kvec_real[[i, 1]]);
423            if kvec_real[[i, 1]] >= 0.0 {
424                s.push_str("    ");
425            } else {
426                s.push_str("   ");
427            }
428            s.push_str(&bb);
429            let cc: String = format!("{:.6}", N_L[[i]]);
430            if N_L[[i]] >= 0.0 {
431                s.push_str("    ");
432            } else {
433                s.push_str("   ");
434            }
435            s.push_str(&cc);
436            let cc: String = format!("{:.6}", N_R[[i]]);
437            if N_R[[i]] >= 0.0 {
438                s.push_str("    ");
439            } else {
440                s.push_str("   ");
441            }
442            let cc: String = format!("{:.6}", N_B[[i]]);
443            if N_B[[i]] >= 0.0 {
444                s.push_str("    ");
445            } else {
446                s.push_str("   ");
447            }
448            s.push_str(&cc);
449            s.push_str("\n");
450        }
451        writer.write_all(s.as_bytes()).unwrap();
452        let _ = file;
453
454        let width: usize = kmesh[[0]];
455        let height: usize = kmesh[[1]];
456
457        let N_L: Array2<f64> =
458            Array2::from_shape_vec((height, width), N_L.to_vec()).expect("Shape error");
459        let N_L = N_L.reversed_axes(); // 转置操作
460        let N_L = N_L.iter().cloned().collect::<Vec<f64>>();
461        let N_R: Array2<f64> =
462            Array2::from_shape_vec((height, width), N_R.to_vec()).expect("Shape error");
463        let N_R = N_R.reversed_axes(); // 转置操作
464        let N_R = N_R.iter().cloned().collect::<Vec<f64>>();
465        let N_B: Array2<f64> =
466            Array2::from_shape_vec((height, width), N_B.to_vec()).expect("Shape error");
467        let N_B = N_B.reversed_axes(); // 转置操作
468        let N_B = N_B.iter().cloned().collect::<Vec<f64>>();
469
470        //接下来我们绘制表面态
471        let mut fg = Figure::new();
472        let mut heatmap_data = N_L;
473        let axes = fg.axes2d();
474        //axes.set_palette(RAINBOW);
475        axes.set_palette(Custom(&[
476            (-1.0, 0.0, 0.0, 0.0),
477            (-0.9, 65.0 / 255.0, 9.0 / 255.0, 103.0 / 255.0),
478            (0.0, 147.0 / 255.0, 37.0 / 255.0, 103.0 / 255.0),
479            (0.2, 220.0 / 255.0, 80.0 / 255.0, 57.0 / 255.0),
480            (1.0, 252.0 / 255.0, 254.0 / 255.0, 164.0 / 255.0),
481        ]));
482        axes.image(
483            heatmap_data.iter(),
484            width,
485            height,
486            Some((0.0, 0.0, 1.0, 1.0)),
487            &[],
488        );
489        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
490        let axes = axes.set_y_range(Fix(0.0), Fix(1.0));
491        let axes = axes.set_aspect_ratio(Fix(1.0));
492        axes.set_x_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
493        axes.set_y_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
494        axes.set_cb_ticks_custom(
495            [
496                Major(-10.0, Fix("low")),
497                Major(0.0, Fix("0")),
498                Major(10.0, Fix("high")),
499            ]
500            .into_iter(),
501            &[],
502            &[Font("Times New Roman", 24.0)],
503        );
504        let mut pdfname = String::new();
505        pdfname.push_str(&name.clone());
506        pdfname.push_str("/surf_state_l.pdf");
507        fg.set_terminal("pdfcairo", &pdfname);
508        fg.show().expect("Unable to draw heatmap");
509        let _ = fg;
510
511        let mut fg = Figure::new();
512        let mut heatmap_data = N_R;
513        let axes = fg.axes2d();
514        //axes.set_palette(RAINBOW);
515        axes.set_palette(Custom(&[
516            (-1.0, 0.0, 0.0, 0.0),
517            (-0.9, 65.0 / 255.0, 9.0 / 255.0, 103.0 / 255.0),
518            (0.0, 147.0 / 255.0, 37.0 / 255.0, 103.0 / 255.0),
519            (0.2, 220.0 / 255.0, 80.0 / 255.0, 57.0 / 255.0),
520            (1.0, 252.0 / 255.0, 254.0 / 255.0, 164.0 / 255.0),
521        ]));
522        axes.image(
523            heatmap_data.iter(),
524            width,
525            height,
526            Some((0.0, 0.0, 1.0, 1.0)),
527            &[],
528        );
529        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
530        let axes = axes.set_y_range(Fix(0.0), Fix(1.0));
531        let axes = axes.set_aspect_ratio(Fix(1.0));
532        axes.set_x_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
533        axes.set_y_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
534        axes.set_cb_ticks_custom(
535            [
536                Major(-10.0, Fix("low")),
537                Major(0.0, Fix("0")),
538                Major(10.0, Fix("high")),
539            ]
540            .into_iter(),
541            &[],
542            &[Font("Times New Roman", 24.0)],
543        );
544        let mut pdfname = String::new();
545        pdfname.push_str(&name.clone());
546        pdfname.push_str("/surf_state_r.pdf");
547        fg.set_terminal("pdfcairo", &pdfname);
548        fg.show().expect("Unable to draw heatmap");
549        let _ = fg;
550
551        let mut fg = Figure::new();
552        let mut heatmap_data = N_B;
553        let axes = fg.axes2d();
554        //axes.set_palette(RAINBOW);
555        axes.set_palette(Custom(&[
556            (-1.0, 0.0, 0.0, 0.0),
557            (-0.9, 65.0 / 255.0, 9.0 / 255.0, 103.0 / 255.0),
558            (0.0, 147.0 / 255.0, 37.0 / 255.0, 103.0 / 255.0),
559            (0.2, 220.0 / 255.0, 80.0 / 255.0, 57.0 / 255.0),
560            (1.0, 252.0 / 255.0, 254.0 / 255.0, 164.0 / 255.0),
561        ]));
562        axes.image(
563            heatmap_data.iter(),
564            width,
565            height,
566            Some((0.0, 0.0, 1.0, 1.0)),
567            &[],
568        );
569        let axes = axes.set_x_range(Fix(0.0), Fix(1.0));
570        let axes = axes.set_y_range(Fix(0.0), Fix(1.0));
571        let axes = axes.set_aspect_ratio(Fix(1.0));
572        axes.set_x_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
573        axes.set_y_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
574        axes.set_cb_ticks_custom(
575            [
576                Major(-10.0, Fix("low")),
577                Major(0.0, Fix("0")),
578                Major(10.0, Fix("high")),
579            ]
580            .into_iter(),
581            &[],
582            &[Font("Times New Roman", 24.0)],
583        );
584        let mut pdfname = String::new();
585        pdfname.push_str(&name.clone());
586        pdfname.push_str("/surf_state_b.pdf");
587        fg.set_terminal("pdfcairo", &pdfname);
588        fg.show().expect("Unable to draw heatmap");
589        let _ = fg;
590    }
591    pub fn show_surf_state(
592        &self,
593        name: &str,
594        kpath: &Array2<f64>,
595        label: &Vec<&str>,
596        nk: usize,
597        E_min: f64,
598        E_max: f64,
599        E_n: usize,
600        spin: usize,
601    ) {
602        use std::fs::create_dir_all;
603        use std::io::{BufWriter, Write};
604        create_dir_all(name).expect("can't creat the file");
605        let (kvec, kdist, knode) = self.k_path(kpath, nk).expect("Failed to generate k-path");
606        let Energy = Array1::<f64>::linspace(E_min, E_max, E_n);
607        let (N_L, N_R, N_B) = self.surf_green_path(&kvec, E_min, E_max, E_n, spin);
608        //let N_L=N_L.mapv(|x| if x > 0.0 {x.ln()} else if  x< 0.0 {-x.abs().ln()} else {0.0});
609        //let N_R=N_R.mapv(|x| if x > 0.0 {x.ln()} else if  x< 0.0 {-x.abs().ln()} else {0.0});
610        //let N_B=N_B.mapv(|x| if x > 0.0 {x.ln()} else if  x< 0.0 {-x.abs().ln()} else {0.0});
611        let ((N_L, N_R), N_B) = if spin == 0 {
612            let N_L = N_L.mapv(|x| x.ln());
613            let N_R = N_R.mapv(|x| x.ln());
614            let N_B = N_B.mapv(|x| x.ln());
615            let max = N_L.iter().fold(f64::NEG_INFINITY, |acc, &x| acc.max(x));
616            let min = N_L.iter().fold(f64::INFINITY, |acc, &x| acc.min(x));
617            let N_L = (N_L - min) / (max - min) * 20.0 - 10.0;
618            let max = N_R.iter().fold(f64::NEG_INFINITY, |acc, &x| acc.max(x));
619            let min = N_R.iter().fold(f64::INFINITY, |acc, &x| acc.min(x));
620            let N_R = (N_R - min) / (max - min) * 20.0 - 10.0;
621            let max = N_B.iter().fold(f64::NEG_INFINITY, |acc, &x| acc.max(x));
622            let min = N_B.iter().fold(f64::INFINITY, |acc, &x| acc.min(x));
623            let N_B = (N_B - min) / (max - min) * 20.0 - 10.0;
624            ((N_L, N_R), N_B)
625        } else {
626            ((N_L, N_R), N_B)
627        };
628
629        //绘制 left_dos------------------------
630        let mut left_name: String = String::new();
631        left_name.push_str(&name.clone());
632        left_name.push_str("/dos.surf_l");
633        let mut file = File::create(left_name).expect("Unable to dos.surf_l.dat");
634        let mut writer = BufWriter::new(file);
635        let mut s = String::new();
636        for i in 0..nk {
637            for j in 0..E_n {
638                let aa = format!("{:.6}", kdist[[i]]);
639                s.push_str(&aa);
640                let bb: String = format!("{:.6}", Energy[[j]]);
641                if Energy[[j]] >= 0.0 {
642                    s.push_str("    ");
643                } else {
644                    s.push_str("   ");
645                }
646                s.push_str(&bb);
647                let cc: String = format!("{:.6}", N_L[[i, j]]);
648                if N_L[[i, j]] >= 0.0 {
649                    s.push_str("    ");
650                } else {
651                    s.push_str("   ");
652                }
653                s.push_str(&cc);
654                s.push_str("    ");
655                //writeln!(file,"{}",s);
656                s.push_str("\n");
657            }
658            s.push_str("\n");
659            //writeln!(file,"\n");
660        }
661        writer.write_all(s.as_bytes()).unwrap();
662        let _ = file;
663
664        //绘制右表面态----------------------
665        let mut left_name: String = String::new();
666        left_name.push_str(&name.clone());
667        left_name.push_str("/dos.surf_r");
668        let mut file = File::create(left_name).expect("Unable to dos.surf_l.dat");
669        let mut writer = BufWriter::new(file);
670        let mut s = String::new();
671        for i in 0..nk {
672            for j in 0..E_n {
673                let aa = format!("{:.6}", kdist[[i]]);
674                s.push_str(&aa);
675                let bb: String = format!("{:.6}", Energy[[j]]);
676                if Energy[[j]] >= 0.0 {
677                    s.push_str("    ");
678                } else {
679                    s.push_str("   ");
680                }
681                s.push_str(&bb);
682                let cc: String = format!("{:.6}", N_R[[i, j]]);
683                if N_L[[i, j]] >= 0.0 {
684                    s.push_str("    ");
685                } else {
686                    s.push_str("   ");
687                }
688                s.push_str(&cc);
689                s.push_str("    ");
690                //writeln!(file,"{}",s);
691                s.push_str("\n");
692            }
693            s.push_str("\n");
694            //writeln!(file,"\n");
695        }
696        writer.write_all(s.as_bytes()).unwrap();
697        let _ = file;
698
699        //绘制体态----------------------
700        let mut left_name: String = String::new();
701        left_name.push_str(&name.clone());
702        left_name.push_str("/dos.surf_bulk");
703        let mut file = File::create(left_name).expect("Unable to dos.surf_l.dat");
704        let mut writer = BufWriter::new(file);
705        let mut s = String::new();
706        for i in 0..nk {
707            for j in 0..E_n {
708                let aa = format!("{:.6}", kdist[[i]]);
709                s.push_str(&aa);
710                let bb: String = format!("{:.6}", Energy[[j]]);
711                if Energy[[j]] >= 0.0 {
712                    s.push_str("    ");
713                } else {
714                    s.push_str("   ");
715                }
716                s.push_str(&bb);
717                let cc: String = format!("{:.6}", N_B[[i, j]]);
718                if N_L[[i, j]] >= 0.0 {
719                    s.push_str("    ");
720                } else {
721                    s.push_str("   ");
722                }
723                s.push_str(&cc);
724                s.push_str("    ");
725                //writeln!(file,"{}",s);
726                s.push_str("\n");
727            }
728            s.push_str("\n");
729            //writeln!(file,"\n");
730        }
731        writer.write_all(s.as_bytes()).unwrap();
732        let _ = file;
733
734        //接下来我们绘制表面态
735        let mut fg = Figure::new();
736        let width: usize = nk;
737        let height: usize = E_n;
738        let mut heatmap_data = vec![];
739        for i in 0..height {
740            for j in 0..width {
741                heatmap_data.push(N_L[[j, i]]);
742            }
743        }
744        let axes = fg.axes2d();
745        //axes.set_palette(RAINBOW);
746        axes.set_palette(Custom(&[
747            (-1.0, 0.0, 0.0, 0.0),
748            (-0.9, 65.0 / 255.0, 9.0 / 255.0, 103.0 / 255.0),
749            (0.0, 147.0 / 255.0, 37.0 / 255.0, 103.0 / 255.0),
750            (0.2, 220.0 / 255.0, 80.0 / 255.0, 57.0 / 255.0),
751            (1.0, 252.0 / 255.0, 254.0 / 255.0, 164.0 / 255.0),
752        ]));
753        axes.image(
754            heatmap_data.iter(),
755            width,
756            height,
757            Some((kdist[[0]], E_min, kdist[[nk - 1]], E_max)),
758            &[],
759        );
760        let axes = axes.set_y_range(Fix(E_min), Fix(E_max));
761        let axes = axes.set_x_range(Fix(kdist[[0]]), Fix(kdist[[nk - 1]]));
762        let axes = axes.set_aspect_ratio(Fix(1.0));
763        let mut show_ticks = Vec::new();
764        for i in 0..knode.len() {
765            let A = knode[[i]];
766            let B = label[i];
767            show_ticks.push(Major(A, Fix(B)));
768        }
769        axes.set_x_ticks_custom(
770            show_ticks.into_iter(),
771            &[],
772            &[Font("Times New Roman", 24.0)],
773        );
774        axes.set_y_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
775        //axes.set_cb_ticks(Some((Fix(5.0),0)),&[],&[Font("Times New Roman",24.0)]);
776        axes.set_cb_ticks_custom(
777            [
778                Major(-10.0, Fix("low")),
779                Major(0.0, Fix("0")),
780                Major(10.0, Fix("high")),
781            ]
782            .into_iter(),
783            &[],
784            &[Font("Times New Roman", 24.0)],
785        );
786        let mut pdfname = String::new();
787        pdfname.push_str(&name.clone());
788        pdfname.push_str("/surf_state_l.pdf");
789        fg.set_terminal("pdfcairo", &pdfname);
790        fg.show().expect("Unable to draw heatmap");
791        let _ = fg;
792
793        //接下来我们绘制right表面态
794        let mut fg = Figure::new();
795        let width: usize = nk;
796        let height: usize = E_n;
797        let mut heatmap_data = vec![];
798        for i in 0..height {
799            for j in 0..width {
800                heatmap_data.push(N_R[[j, i]]);
801            }
802        }
803        let axes = fg.axes2d();
804        //axes.set_palette(RAINBOW);
805        axes.set_palette(Custom(&[
806            (-1.0, 0.0, 0.0, 0.0),
807            (-0.9, 65.0 / 255.0, 9.0 / 255.0, 103.0 / 255.0),
808            (0.0, 147.0 / 255.0, 37.0 / 255.0, 103.0 / 255.0),
809            (0.2, 220.0 / 255.0, 80.0 / 255.0, 57.0 / 255.0),
810            (1.0, 252.0 / 255.0, 254.0 / 255.0, 164.0 / 255.0),
811        ]));
812        axes.image(
813            heatmap_data.iter(),
814            width,
815            height,
816            Some((kdist[[0]], E_min, kdist[[nk - 1]], E_max)),
817            &[],
818        );
819        let axes = axes.set_y_range(Fix(E_min), Fix(E_max));
820        let axes = axes.set_x_range(Fix(kdist[[0]]), Fix(kdist[[nk - 1]]));
821        let axes = axes.set_aspect_ratio(Fix(1.0));
822        let mut show_ticks = Vec::new();
823        for i in 0..knode.len() {
824            let A = knode[[i]];
825            let B = label[i];
826            show_ticks.push(Major(A, Fix(B)));
827        }
828        axes.set_x_ticks_custom(
829            show_ticks.into_iter(),
830            &[],
831            &[Font("Times New Roman", 24.0)],
832        );
833        axes.set_y_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
834        //axes.set_cb_ticks(Some((Fix(5.0),0)),&[],&[Font("Times New Roman",24.0)]);
835        axes.set_cb_ticks_custom(
836            [
837                Major(-10.0, Fix("low")),
838                Major(0.0, Fix("0")),
839                Major(10.0, Fix("high")),
840            ]
841            .into_iter(),
842            &[],
843            &[Font("Times New Roman", 24.0)],
844        );
845        let mut pdfname = String::new();
846        pdfname.push_str(&name.clone());
847        pdfname.push_str("/surf_state_r.pdf");
848        fg.set_terminal("pdfcairo", &pdfname);
849        fg.show().expect("Unable to draw heatmap");
850        let _ = fg;
851        //接下来我们绘制bulk表面态
852        let mut fg = Figure::new();
853        let width: usize = nk;
854        let height: usize = E_n;
855        let mut heatmap_data = vec![];
856        for i in 0..height {
857            for j in 0..width {
858                heatmap_data.push(N_B[[j, i]]);
859            }
860        }
861        let axes = fg.axes2d();
862        //axes.set_palette(RAINBOW);
863        axes.set_palette(Custom(&[
864            (-1.0, 0.0, 0.0, 0.0),
865            (-0.9, 65.0 / 255.0, 9.0 / 255.0, 103.0 / 255.0),
866            (0.0, 147.0 / 255.0, 37.0 / 255.0, 103.0 / 255.0),
867            (0.2, 220.0 / 255.0, 80.0 / 255.0, 57.0 / 255.0),
868            (1.0, 252.0 / 255.0, 254.0 / 255.0, 164.0 / 255.0),
869        ]));
870        axes.image(
871            heatmap_data.iter(),
872            width,
873            height,
874            Some((kdist[[0]], E_min, kdist[[nk - 1]], E_max)),
875            &[],
876        );
877        let axes = axes.set_y_range(Fix(E_min), Fix(E_max));
878        let axes = axes.set_x_range(Fix(kdist[[0]]), Fix(kdist[[nk - 1]]));
879        let axes = axes.set_aspect_ratio(Fix(1.0));
880        let mut show_ticks = Vec::new();
881        for i in 0..knode.len() {
882            let A = knode[[i]];
883            let B = label[i];
884            show_ticks.push(Major(A, Fix(B)));
885        }
886        axes.set_x_ticks_custom(
887            show_ticks.into_iter(),
888            &[],
889            &[Font("Times New Roman", 24.0)],
890        );
891        axes.set_y_ticks(Some((Auto, 0)), &[], &[Font("Times New Roman", 24.0)]);
892        //axes.set_cb_ticks(Some((Fix(5.0),0)),&[],&[Font("Times New Roman",24.0)]);
893        axes.set_cb_ticks_custom(
894            [
895                Major(-10.0, Fix("low")),
896                Major(0.0, Fix("0")),
897                Major(10.0, Fix("high")),
898            ]
899            .into_iter(),
900            &[],
901            &[Font("Times New Roman", 24.0)],
902        );
903        let mut pdfname = String::new();
904        pdfname.push_str(&name.clone());
905        pdfname.push_str("/surf_state_b.pdf");
906        fg.set_terminal("pdfcairo", &pdfname);
907        fg.show().expect("Unable to draw heatmap");
908        let _ = fg;
909    }
910}