From c4fe6c61c65e165458cd6fe5e9153f452053ec79 Mon Sep 17 00:00:00 2001
From: HeFuxiang94 <78629482+HeFuxiang94@users.noreply.github.com>
Date: Mon, 20 Feb 2023 20:36:07 +0800
Subject: [PATCH 1/2] Add the Dipole and post-processing script necessary to
 calculate the spectrum using TDDFT.

---
 source/module_io/write_dipole.cpp |   9 +--
 tools/ABACUS_dipole_plot.py       | 102 ++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 4 deletions(-)
 create mode 100644 tools/ABACUS_dipole_plot.py

diff --git a/source/module_io/write_dipole.cpp b/source/module_io/write_dipole.cpp
index 0b2fc11ff3..ff78808e13 100644
--- a/source/module_io/write_dipole.cpp
+++ b/source/module_io/write_dipole.cpp
@@ -196,7 +196,7 @@ void ModuleIO::write_dipole(const double *rho_save,
 
         ofs << istep << " " << dipole_elec_x << " " << dipole_elec_y << " " << dipole_elec_z << std::endl;
 
-        /*
+        
         double dipole_ion_x = 0.0, dipole_ion_y = 0.0, dipole_ion_z = 0.0, dipole_sum = 0.0;
         if (GlobalC::ucell.ntype == 1)
         {
@@ -265,6 +265,7 @@ void ModuleIO::write_dipole(const double *rho_save,
         {
             std::cout << "atom ntype is too large!" << std::endl;
         }
+/*
         for(int it=1; it<(GlobalC::ucell.ntype); it++)
                 {
                     for(int ia=0; ia<GlobalC::ucell.atoms[it].na; ia++)
@@ -289,7 +290,7 @@ void ModuleIO::write_dipole(const double *rho_save,
 
                     }
                 }
-        
+*/     
 
         std::cout << std::setprecision(8) << "dipole_ion_x: " << dipole_ion_x << std::endl;
         std::cout << std::setprecision(8) << "dipole_ion_y: " << dipole_ion_y << std::endl;
@@ -302,8 +303,8 @@ void ModuleIO::write_dipole(const double *rho_save,
         std::cout << std::setprecision(8) << "dipole_x: " << dipole_x << std::endl;
         std::cout << std::setprecision(8) << "dipole_y: " << dipole_y << std::endl;
         std::cout << std::setprecision(8) << "dipole_z: " << dipole_z << std::endl;
-        dipole_sum = sqrt(dipole_x * dipole_x + dipole_y * dipole_y + dipole_z * dipole_z);
-        std::cout << std::setprecision(8) << "dipole_sum: " << dipole_sum << std::endl;*/
+        //dipole_sum = sqrt(dipole_x * dipole_x + dipole_y * dipole_y + dipole_z * dipole_z);
+        //std::cout << std::setprecision(8) << "dipole_sum: " << dipole_sum << std::endl;
     }
     MPI_Barrier(MPI_COMM_WORLD);
 #endif
diff --git a/tools/ABACUS_dipole_plot.py b/tools/ABACUS_dipole_plot.py
new file mode 100644
index 0000000000..72f9b1f9d4
--- /dev/null
+++ b/tools/ABACUS_dipole_plot.py
@@ -0,0 +1,102 @@
+from re import split
+import numpy as np
+import matplotlib.pyplot as plt
+from ase.units import Bohr
+from math import pi
+from scipy.constants import c
+from os import listdir
+from os import mkdir
+from shutil import move
+from shutil import copyfile
+from os import path
+def Sp_plot(t_name,out,file,move=0,N_t=-1,dt=0.0034,direction=-1):
+    Dipole=dipole_read(t_name,out,file,direction=direction)
+    if(N_t==-1):
+        N_t=Dipole.shape[1]
+        N_t-=move
+    shift_omega_step=1
+    lambda_x=(c/np.linspace(shift_omega_step,N_t-1,N_t-shift_omega_step)*dt*N_t*1e-6)
+    fig,ax=plt.subplots()
+    if(direction==1 or direction==-1):
+        alpha_x=np.fft.fft(Dipole[0][move:N_t])
+        S_x=np.abs(alpha_x.imag)
+        ax.plot(lambda_x,S_x[shift_omega_step:],label='X')
+    if(direction==2 or direction==-1):
+        alpha_y=np.fft.fft(Dipole[1][move:N_t])
+        S_y=np.abs(alpha_y.imag)
+        ax.plot(lambda_x,S_y[shift_omega_step:],label='Y')
+    if(direction==3 or direction==-1):
+        alpha_z=np.fft.fft(Dipole[2][move:N_t])
+        S_z=np.abs(alpha_z.imag)
+        ax.plot(lambda_x,S_x[shift_omega_step:],label='Z') 
+    ax.set_ylabel('Strength')
+    ax.set_xlabel('frequency(eV)')
+    ax.set_title("Spectra")
+    ax.set_xlim(80,200)
+    plt.legend()
+    plt.savefig('./'+out+'/'+file+'spec.png')
+
+def Sp_field_plot():
+    pass
+def dipole_read(total_name,out,file,dt=0.0034,direction=-1):
+    #Initial
+    list_d=[[] for i in range(3)]
+    #N_t=nstep-move_step                 #step used in fft
+    #Read in cycle
+    m=0                                 #Dimension label
+    line=1
+    with open(total_name,mode='r')as f1:
+        while(line!=''):
+            while(m<3):
+                line=f1.readline()
+                if(line==''):break
+                tmp=line.split()
+                if (m==0):
+                    list_d[0].append(float(tmp[1]))
+                elif (m==1):
+                    list_d[1].append(float(tmp[1]))
+                elif (m==2):
+                    list_d[2].append(float(tmp[1]))
+                m+=1
+            m=0
+    #Plot part
+    Dipole=np.array(list_d)
+    N_t=Dipole.shape[1]
+    Ti=np.linspace(0,(N_t-1)*dt,N_t)
+    fig,ax=plt.subplots()
+    if(direction==1 or direction==-1):
+        ax.plot(Ti,Dipole[0],label='X')
+    if(direction==2 or direction==-1):
+        ax.plot(Ti,Dipole[1],label='Y')
+    if(direction==3 or direction==-1):
+        ax.plot(Ti,Dipole[2],label='Z')
+    ax.set_ylabel('Strength')
+    ax.set_xlabel('time(fs)')
+    ax.set_title('Dipole')
+    plt.legend()
+    plt.savefig('./'+out+'/'+file+'dipole.png')
+    return Dipole
+#reading para
+def autoanlysis(filename,dir,out):
+    list_dir=listdir("./"+dir+"/")
+    count=0
+    print("total file: "+str(len(list_dir)))
+    #main cycle
+    for file in list_dir:
+        if(filename in file):
+            count+=1
+            t_name="./"+dir+"/"+file
+            Sp_plot(t_name,out,file)
+
+#main part
+filename='dipole.dat'
+dir='dipole_nomove'
+out='PIC'
+file='C2H2_nomove_X_dipole.dat'
+if(not path.exists('./'+out)):
+    mkdir(out)
+#direction='X'
+#autoanlysis(filename,dir,out)
+#single plot
+t_name="./"+dir+"/"+file
+Sp_plot(t_name,out,file,direction=1)
\ No newline at end of file

From a0f8a0c459ae70f110b71d845460e8ad3a4bbb3f Mon Sep 17 00:00:00 2001
From: HeFuxiang94 <78629482+HeFuxiang94@users.noreply.github.com>
Date: Wed, 22 Feb 2023 11:17:56 +0800
Subject: [PATCH 2/2] fix td_val test bug.

---
 source/module_io/write_dipole.cpp | 36 +++++++++++++++----------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/source/module_io/write_dipole.cpp b/source/module_io/write_dipole.cpp
index ff78808e13..ea97c93870 100644
--- a/source/module_io/write_dipole.cpp
+++ b/source/module_io/write_dipole.cpp
@@ -203,11 +203,11 @@ void ModuleIO::write_dipole(const double *rho_save,
             for (int ia = 0; ia < GlobalC::ucell.atoms[0].na; ia++)
             {
                 dipole_ion_x += GlobalC::ucell.atoms[0].taud[ia].x * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
                 dipole_ion_y += GlobalC::ucell.atoms[0].taud[ia].y * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
                 dipole_ion_z += GlobalC::ucell.atoms[0].taud[ia].z * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
             }
         }
         else if (GlobalC::ucell.ntype == 2)
@@ -215,20 +215,20 @@ void ModuleIO::write_dipole(const double *rho_save,
             for (int ia = 0; ia < GlobalC::ucell.atoms[0].na; ia++)
             {
                 dipole_ion_x += GlobalC::ucell.atoms[0].taud[ia].x * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
                 dipole_ion_y += GlobalC::ucell.atoms[0].taud[ia].y * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
                 dipole_ion_z += GlobalC::ucell.atoms[0].taud[ia].z * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
             }
             for (int ia = 0; ia < GlobalC::ucell.atoms[1].na; ia++)
             {
                 dipole_ion_x += GlobalC::ucell.atoms[1].taud[ia].x * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_02;
+                                * INPUT.td_val_elec_02;
                 dipole_ion_y += GlobalC::ucell.atoms[1].taud[ia].y * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_02;
+                                * INPUT.td_val_elec_02;
                 dipole_ion_z += GlobalC::ucell.atoms[1].taud[ia].z * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_02;
+                                * INPUT.td_val_elec_02;
             }
         }
         else if (GlobalC::ucell.ntype == 3)
@@ -236,29 +236,29 @@ void ModuleIO::write_dipole(const double *rho_save,
             for (int ia = 0; ia < GlobalC::ucell.atoms[0].na; ia++)
             {
                 dipole_ion_x += GlobalC::ucell.atoms[0].taud[ia].x * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
                 dipole_ion_y += GlobalC::ucell.atoms[0].taud[ia].y * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
                 dipole_ion_z += GlobalC::ucell.atoms[0].taud[ia].z * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_01;
+                                * INPUT.td_val_elec_01;
             }
             for (int ia = 0; ia < GlobalC::ucell.atoms[1].na; ia++)
             {
                 dipole_ion_x += GlobalC::ucell.atoms[1].taud[ia].x * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_02;
+                                * INPUT.td_val_elec_02;
                 dipole_ion_y += GlobalC::ucell.atoms[1].taud[ia].y * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_02;
+                                * INPUT.td_val_elec_02;
                 dipole_ion_z += GlobalC::ucell.atoms[1].taud[ia].z * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_02;
+                                * INPUT.td_val_elec_02;
             }
             for (int ia = 0; ia < GlobalC::ucell.atoms[2].na; ia++)
             {
                 dipole_ion_x += GlobalC::ucell.atoms[2].taud[ia].x * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_03;
+                                * INPUT.td_val_elec_03;
                 dipole_ion_y += GlobalC::ucell.atoms[2].taud[ia].y * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_03;
+                                * INPUT.td_val_elec_03;
                 dipole_ion_z += GlobalC::ucell.atoms[2].taud[ia].z * GlobalC::ucell.lat0 * 0.529177
-                                * ELEC_evolve::td_val_elec_03;
+                                * INPUT.td_val_elec_03;
             }
         }
         else