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

Rustb/
surfgreen.rs
