From 9ab7fd704ba6bb13550f2c7663e297013cf3b80e Mon Sep 17 00:00:00 2001 From: Fisherd99 Date: Fri, 31 Jul 2026 14:04:05 +0800 Subject: [PATCH 1/2] 1. feat: add BSE module 2. feat: add ExcitonPlotter 3. add input parameter `rpa_outdir` and `rpa_out_vel` --- docs/advanced/input_files/input-main.md | 201 +++- .../module_external/lapack_connector.h | 21 +- .../module_external/scalapack_connector.h | 72 +- source/source_base/parallel_2d.cpp | 14 +- source/source_base/parallel_2d.h | 10 +- .../source_io/module_ctrl/ctrl_scf_lcao.cpp | 2 + .../module_parameter/input_parameter.h | 23 +- .../read_input_item_output.cpp | 29 +- .../read_input_item_tddft.cpp | 124 ++ source/source_io/test/read_input_ptest.cpp | 24 +- source/source_lcao/module_lr/CMakeLists.txt | 14 + .../module_lr/ao_to_mo_transformer/ao_to_mo.h | 14 +- .../ao_to_mo_parallel.cpp | 30 +- .../ao_to_mo_transformer/ao_to_mo_serial.cpp | 42 +- .../test/ao_to_mo_test.cpp | 54 +- .../source_lcao/module_lr/bse/CMakeLists.txt | 5 + source/source_lcao/module_lr/bse/bse_util.cpp | 151 +++ source/source_lcao/module_lr/bse/bse_util.h | 96 ++ .../module_lr/bse/esolver_bse_lcao.cpp | 731 ++++++++++++ .../module_lr/bse/esolver_bse_lcao.h | 57 + .../source_lcao/module_lr/bse/hamilt_bse.cpp | 603 ++++++++++ source/source_lcao/module_lr/bse/hamilt_bse.h | 120 ++ .../module_lr/bse/hamilt_bse_solver.cpp | 331 ++++++ .../module_lr/bse/hamilt_bse_solver.h | 53 + .../module_lr/bse/hamilt_bse_solver.hpp | 62 + .../source_lcao/module_lr/bse/molecular_lri.h | 164 +++ .../module_lr/bse/molecular_lri.hpp | 289 +++++ .../module_lr/bse/molecular_lri_comm.hpp | 276 +++++ .../module_lr/bse/test/CMakeLists.txt | 6 + .../bse/test/hamilt_bse_solver_test.cpp | 249 ++++ .../source_lcao/module_lr/dm_trans/dm_trans.h | 12 +- .../module_lr/dm_trans/dm_trans_parallel.cpp | 26 +- .../module_lr/dm_trans/dm_trans_serial.cpp | 32 +- .../module_lr/dm_trans/test/dm_trans_test.cpp | 38 +- .../source_lcao/module_lr/hamilt_casida.cpp | 8 +- source/source_lcao/module_lr/hamilt_casida.h | 82 +- source/source_lcao/module_lr/hamilt_ulr.hpp | 62 +- source/source_lcao/module_lr/hsolver_lrtd.hpp | 13 +- source/source_lcao/module_lr/lr_spectrum.cpp | 260 +++- source/source_lcao/module_lr/lr_spectrum.h | 47 +- .../module_lr/lr_spectrum_velocity.cpp | 219 +++- .../operator_casida/operator_lr_diag.h | 2 + .../operator_casida/operator_lr_exx.cpp | 42 +- .../operator_casida/operator_lr_exx.h | 14 +- .../operator_casida/operator_lr_hxc.cpp | 11 +- .../operator_casida/operator_lr_hxc.h | 3 +- .../ri_benchmark/operator_ri_hartree.h | 10 +- .../module_lr/ri_benchmark/ri_benchmark.h | 19 +- .../module_lr/ri_benchmark/ri_benchmark.hpp | 122 +- .../module_lr/utils/exciton_plotter.cpp | 1042 +++++++++++++++++ .../module_lr/utils/exciton_plotter.h | 265 +++++ source/source_lcao/module_lr/utils/lr_io.cpp | 948 +++++++++++++++ source/source_lcao/module_lr/utils/lr_io.h | 122 ++ source/source_lcao/module_lr/utils/lr_util.h | 64 +- .../source_lcao/module_lr/utils/lr_util.hpp | 463 +++++++- .../module_lr/utils/lr_util_hcontainer.h | 5 +- .../module_lr/utils/lr_util_print.h | 8 + source/source_lcao/module_lr/utils/mo_type.h | 43 + .../module_lr/utils/spectrum_mo.hpp | 208 ++++ source/source_lcao/module_ri/LRI_CV_Tools.h | 2 + source/source_lcao/module_ri/RPA_LRI.h | 8 +- source/source_lcao/module_ri/RPA_LRI.hpp | 80 +- source/source_lcao/module_ri/write_ri_cv.hpp | 68 +- tests/08_EXX/51_GO_LR/KPT | 4 + tests/08_EXX/51_GO_LR/result.ref | 6 +- tests/08_EXX/52_GO_LR_PBE/KPT | 4 + tests/08_EXX/52_GO_LR_PBE/result.ref | 6 +- tests/08_EXX/53_GO_LR_HF/KPT | 4 + tests/08_EXX/53_GO_LR_HF/result.ref | 7 +- tests/08_EXX/54_GO_ULR_HF/KPT | 4 + tests/08_EXX/54_GO_ULR_HF/result.ref | 6 +- tests/08_EXX/55_KP_LR/result.ref | 7 +- tests/08_EXX/56_GO_RPA_OUTPUT/result.ref | 31 +- tests/08_EXX/57_KP_RPA_SHRINK/result.ref | 47 +- tests/integrate/tools/catch_properties.sh | 12 +- .../plot-tools/plot_cond_silce.py | 385 ++++++ 76 files changed, 8099 insertions(+), 639 deletions(-) create mode 100644 source/source_lcao/module_lr/bse/CMakeLists.txt create mode 100644 source/source_lcao/module_lr/bse/bse_util.cpp create mode 100644 source/source_lcao/module_lr/bse/bse_util.h create mode 100644 source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp create mode 100644 source/source_lcao/module_lr/bse/esolver_bse_lcao.h create mode 100644 source/source_lcao/module_lr/bse/hamilt_bse.cpp create mode 100644 source/source_lcao/module_lr/bse/hamilt_bse.h create mode 100644 source/source_lcao/module_lr/bse/hamilt_bse_solver.cpp create mode 100644 source/source_lcao/module_lr/bse/hamilt_bse_solver.h create mode 100644 source/source_lcao/module_lr/bse/hamilt_bse_solver.hpp create mode 100644 source/source_lcao/module_lr/bse/molecular_lri.h create mode 100644 source/source_lcao/module_lr/bse/molecular_lri.hpp create mode 100644 source/source_lcao/module_lr/bse/molecular_lri_comm.hpp create mode 100644 source/source_lcao/module_lr/bse/test/CMakeLists.txt create mode 100644 source/source_lcao/module_lr/bse/test/hamilt_bse_solver_test.cpp create mode 100644 source/source_lcao/module_lr/utils/exciton_plotter.cpp create mode 100644 source/source_lcao/module_lr/utils/exciton_plotter.h create mode 100644 source/source_lcao/module_lr/utils/lr_io.cpp create mode 100644 source/source_lcao/module_lr/utils/lr_io.h create mode 100644 source/source_lcao/module_lr/utils/mo_type.h create mode 100644 source/source_lcao/module_lr/utils/spectrum_mo.hpp create mode 100644 tests/08_EXX/51_GO_LR/KPT create mode 100644 tests/08_EXX/52_GO_LR_PBE/KPT create mode 100644 tests/08_EXX/53_GO_LR_HF/KPT create mode 100644 tests/08_EXX/54_GO_ULR_HF/KPT create mode 100644 tools/02_postprocessing/plot-tools/plot_cond_silce.py diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md index 17615f6174..012d06a58e 100644 --- a/docs/advanced/input_files/input-main.md +++ b/docs/advanced/input_files/input-main.md @@ -201,6 +201,8 @@ - [out\_element\_info](#out_element_info) - [restart\_save](#restart_save) - [rpa](#rpa) + - [rpa_out_vel](#rpa_out_vel) + - [rpa_outdir](#rpa_outdir) - [out\_pchg](#out_pchg) - [out\_wfc\_norm](#out_wfc_norm) - [out\_wfc\_re\_im](#out_wfc_re_im) @@ -540,9 +542,6 @@ - [pexsi\_elec\_thr](#pexsi_elec_thr) - [pexsi\_zero\_thr](#pexsi_zero_thr) - [Linear Response TDDFT](#linear-response-tddft) - - [ri\_hartree\_benchmark](#ri_hartree_benchmark) - - [aims\_nbasis](#aims_nbasis) - - [Linear Response TDDFT (Under Development Feature)](#linear-response-tddft-under-development-feature) - [xc\_kernel](#xc_kernel) - [lr\_init\_xc\_kernel](#lr_init_xc_kernel) - [lr\_solver](#lr_solver) @@ -555,6 +554,23 @@ - [out\_wfc\_lr](#out_wfc_lr) - [abs\_gauge](#abs_gauge) - [abs\_broadening](#abs_broadening) + - [bse\_tda](#bse_tda) + - [bse\_spin\_types](#bse_spin_types) + - [bse\_mem\_save](#bse_mem_save) + - [bse\_ri\_hartree](#bse_ri_hartree) + - [bse\_use\_fine\_kgrid](#bse_use_fine_kgrid) + - [out\_bse\_ab](#out_bse_ab) + - [bse\_continue](#bse_continue) + - [plot\_istate](#plot_istate) + - [exciton\_plot\_type](#exciton_plot_type) + - [exciton\_plot\_format](#exciton_plot_format) + - [exciton\_fixed\_coordinate](#exciton_fixed_coordinate) + - [exciton\_slice\_plane](#exciton_slice_plane) + - [exciton\_slice\_pos](#exciton_slice_pos) + - [exciton\_slice\_npoints](#exciton_slice_npoints) + - [exciton\_slice\_scale](#exciton_slice_scale) + - [ri\_hartree\_benchmark](#ri_hartree_benchmark) + - [aims\_nbasis](#aims_nbasis) - [Reduced Density Matrix Functional Theory](#reduced-density-matrix-functional-theory) - [rdmft](#rdmft) - [rdmft\_power\_alpha](#rdmft_power_alpha) @@ -2304,11 +2320,23 @@ ### rpa - **Type**: Boolean -- **Description**: Generate output files used in rpa calculations. +- **Description**: Generate output files used in rpa calculations for LibRPA package. > Note: If symmetry is set to 1, additional files containing the necessary information for exploiting symmetry in the subsequent rpa calculation will be output: irreducible_sector.txt, symrot_k.txt and symrot_R.txt. - **Default**: False +### rpa_out_vel + +- **Type**: Boolean +- **Description**: Whether to output velocity matrix for LibRPA package. +- **Default**: False + +### rpa_outdir + +- **Type**: String +- **Description**: Output directory for LibRPA package. +- **Default**: "./OUT_rpa/" + ### out_pchg - **Type**: String @@ -4768,27 +4796,13 @@ ## Linear Response TDDFT -### ri_hartree_benchmark - -- **Type**: String -- **Description**: Whether to use the RI approximation for the Hartree term in LR-TDDFT for benchmark (with FHI-aims/ABACUS read-in style) -- **Default**: none - -### aims_nbasis - -- **Type**: A number(ntype) of Integers -- **Availability**: *ri_hartree_benchmark = aims* -- **Description**: Atomic basis set size for each atom type (with the same order as in STRU) in FHI-aims. -- **Default**: {} (empty list, where ABACUS use its own basis set size) - -[back to top](#full-list-of-input-keywords) - -## Linear Response TDDFT (Under Development Feature) +These parameters are used to solve the excited states using. e.g. LR-TDDFT or Bethe-Salpeter Equation. ### xc_kernel - **Type**: String -- **Description**: The exchange-correlation kernel used in the calculation. Currently supported: RPA, LDA, PBE, HSE, HF. +- **Description**: The exchange-correlation kernel used in the calculation. +Currently supported: `RPA`, `LDA`, `PBE`, `HSE`, `HF`, `BSE`. - **Default**: LDA ### lr_init_xc_kernel @@ -4803,10 +4817,12 @@ ### lr_solver - **Type**: String -- **Description**: The method to solve the Casida equation in LR-TDDFT under Tamm-Dancoff approximation (TDA). - - dav/dav_subspace/cg: Construct and diagonalize the Hamiltonian matrix iteratively with Davidson/Non-ortho-Davidson/CG algorithm. - - lapack: Construct the full matrix and directly diagonalize with LAPACK. - - spectrum: Calculate absorption spectrum only without solving Casida equation. +- **Description**: The method to solve the Casida equation in LR-TDDFT. + - `dav`/`dav_subspace`/ `cg`: Construct and diagonalize the Hamiltonian matrix iteratively with Davidson/Non-ortho-Davidson/CG algorithm. Tamm-Dancoff approximation (TDA) only. + - `lapack`: Construct the matrix and directly diagonalize with LAPACK. TDA only. + - `elpa`: Construct the resonant (and anti-resonant) matrix and diagonalize with ELPA. + - `spectrum`: Calculate absorption spectrum only without solving Casida equation. The `OUT.${suffix}/` directory should contain the files for LR-TDDFT eigenstates and eigenvalues, i.e. `Excitation_Energy.dat` and `Excitation_Amplitude_${processor_rank}.dat` output by setting `out_wfc_lr` to true. + - `plot`: Plot the exciton wave function, should identify `plot_istate`. - **Default**: dav ### lr_thr @@ -4858,8 +4874,8 @@ ### abs_gauge - **Type**: String -- **Description**: Whether to use length or velocity gauge to calculate the absorption spectrum in LR-TDDFT. -- **Default**: length +- **Description**: Whether to use `velocity` or `length` formulation to calculate the absorption spectrum. `length` is only suitable for non-periodic system. +- **Default**: velocity ### abs_broadening @@ -4867,10 +4883,141 @@ - **Description**: The broadening factor for the absorption spectrum calculation. - **Default**: 0.01 + +### bse_tda + +- **Type**: String +- **Description**: Whether Tamm-Dancoff Approximation is used (can be 'tda', 'full' or 'both'). +- **Default**: tda + +### bse_spin_types + +- **Type**: Vector of String +- **Description**: spin types for close-shell case to be calculated in one task (can be 'singlet', 'triplet', and for test 'rpa', 'ipa'). +- **Defalut**: \{singlet, triplet\} + +### bse_mem_save + +- **Type**: Boolean +- **Description**: Whether to save memory by adding V and W to BSE matrix directly. This option is useful when out-of-memory occurs. If this option is on, `bse_ri_hartree` will be on and `bse_continue` will be off automatically. +- **Default**: false + +### bse_ri_hartree + +- **Type**: Boolean +- **Description**: Whether to use RI approximation for Hartree term in BSE. +- **Default**: true + +### bse_use_fine_kgrid + +- **Type**: Boolean +- **Description**: Whether to use a finer k-grid for BSE. If you want to turn it on, file `band_kpath_info`, `band_KS_eigenvector_k_{index}.txt`, `KS_band_spin_{index}.txt` and `GW_band_spin_{index}.txt` should be prepared. +- **Default**: false + +### out_bse_ab + +- **Type**: Boolean +- **Description**: Whether to output the AB matrix to file. +- **Default**: false + +### bse_continue + +- **Type**: Integer +- **Description**: Which step to continue from previous BSE calculation. + - 0: new; + - 1: continue from A_V; + - 2: continue from A_V and A_W; + - 3: continue from A_V, A_W and B_V; + - 4: continue from A_V, A_W, B_V and B_W +- **Default**: 0 + +### plot_istate + +- **Type**: Integer +- **Description**: The index of excited state to be plotted (starting from 0) +- **Default**: 0 + +### exciton_plot_type + +- **Type**: String +- **Description**: The exciton density represented when `lr_solver` is `plot`. + - `average`: Integrate out the other particle coordinate and plot the average electron and hole densities. + - `conditional`: Fix one particle at the Cartesian position specified by the exciton fixed-coordinate parameters and plot the density of the other particle. +- **Default**: average + +### exciton_plot_format + +- **Type**: String +- **Description**: The output format for exciton-density plotting. + - `cube`: Write three-dimensional density files in Gaussian cube format. + - `slice`: Write two-dimensional cross-section data files. + - `both`: Write both cube and slice files. + - The default value is `cube`. +- **Default**: cube + +### exciton_fixed_coordinate + +- **Type**: Vector of Real (6 values) +- **Description**: Cartesian coordinates in Bohr used by `exciton_plot_type = conditional`, in the order `hole_x hole_y hole_z electron_x electron_y electron_z`. +- **Default**: 0.0 0.0 0.0 0.0 0.0 0.0 +- **Unit**: Bohr + +### exciton_slice_plane + +- **Type**: String +- **Description**: Pair of lattice-vector directions spanning the cross section. + - `ab`: Plane spanned by lattice vectors **a** and **b**. + - `bc`: Plane spanned by lattice vectors **b** and **c**. + - `ca`: Plane spanned by lattice vectors **c** and **a**. +- **Default**: ab + +### exciton_slice_pos + +- **Type**: Real +- **Description**: Offset of the slice along the direction of the remaining lattice vector: **c** for `ab`, **a** for `bc`, and **b** for `ca`. +- **Default**: 0.0 +- **Unit**: Bohr + +### exciton_slice_npoints + +- **Type**: Integer +- **Description**: Target in-plane grid resolution. The final number of points is adjusted to use a uniform number of points per primitive cell over the plotted BvK range and includes both endpoints. +- **Default**: 200 + +### exciton_slice_scale + +- **Type**: Real +- **Description**: Scale factor for the in-plane range relative to the BvK supercell. Values smaller than 1.0 are treated as 1.0. +- **Default**: 1.3 + +### ri_hartree_benchmark + +- **Type**: String +- **Description**: Whether to use the localized resolution-of-identity (LRI) approximation for the **Hartree** term of kernel in the $A$ matrix of LR-TDDFT for benchmark (with FHI-aims or another ABACUS calculation). + - `aims`: The `read_file_dir` directory should contain the FHI-aims output files: RI-LVL tensors `Cs_data_0.txt` and `coulomb_mat_0.txt`, and KS eigenstates from FHI-aims: `band_out`and `KS_eigenvectors.out`. The Casida equation will be constructed under FHI-aims' KS eigenpairs. + - LRI tensor files (`Cs_data_0.txt` and `coulomb_mat_0.txt`)and Kohn-Sham eigenvalues (`bands_out`): run FHI-aims with periodic boundary conditions and with `total_energy_method rpa` and `output librpa`. + - Kohn-Sham eigenstates under aims NAOs (`KS_eigenvectors.out`): run FHI-aims with `output eigenvectors`. + - If the number of atomic orbitals of any atom type in FHI-aims is different from that in ABACUS, the `aims_nbasis` should be set. + - `abacus`: The `read_file_dir` directory should contain the RI-LVL tensors `Cs` and `Vs` (written by setting `out_ri_cv` to 1). The Casida equation will be constructed under ABACUS' KS eigenpairs, with the only difference that the Hartree term is constructed with RI approximation. + - `aims-librpa`: The current directory where you are running task should contain the the LRI tensors `Cs_data_${processor_rank}.txt`, `coulomb_mat${processor_rank}.txt`, and KS eigenstates `band_out`, `KS_eigenvector_${processor_rank}.dat`. All these files can get from FHI-aims task with periodic boundary conditions and with `total_energy_method rpa` and `output librpa`. The `aims_nbasis` should also 0be set. + - `abacus-librpa`: The current directory where you are running task should contain the the LRI tensors `Cs_data_${processor_rank}.txt`, `coulomb_mat_${processor_rank}.txt`, and KS eigenstates `band_out`, `KS_eigenvector_${kpoint}.dat`. All these files can get from ABACUS task with `RPA 1`. + - `none`: Construct the Hartree term by Poisson equation and grid integration as usual. +- **Default**: none + +### aims_nbasis + +- **Type**: A number(ntype) of Integers +- **Availability**: `ri_hartree_benchmark` = `aims` or `aims-librpa` +- **Description**: Atomic basis set size for each atom type (with the same order as in `STRU`) in FHI-aims. +- **Default**: {} (empty list, where ABACUS use its own basis set size) + [back to top](#full-list-of-input-keywords) ## Reduced Density Matrix Functional Theory +Ab-initio methods and the xc-functional parameters used in RDMFT. +The physical quantities that RDMFT temporarily expects to output are the kinetic energy, total energy, and 1-RDM of the system in the ground state, etc. + ### rdmft - **Type**: Boolean diff --git a/source/source_base/module_external/lapack_connector.h b/source/source_base/module_external/lapack_connector.h index e5e576c506..c71e926eb8 100644 --- a/source/source_base/module_external/lapack_connector.h +++ b/source/source_base/module_external/lapack_connector.h @@ -318,6 +318,10 @@ void dsysv_(const char* uplo, const int* n, const int* nrhs, double* b, const int* ldb, double* work, const int* lwork, int* info); + +// === calculate matrix norm === +double dlange_(const char* norm, const int* m, const int* n, const double* A, const int* lda, double* work); +double zlange_(const char* norm, const int* m, const int* n, const std::complex* A, const int* lda, double* work); } // extern "C" #ifdef GATHER_INFO @@ -598,5 +602,18 @@ namespace LapackConnector const char trans_changed = change_trans_NC(trans); cherk_(&uplo_changed, &trans_changed, &n, &k, &alpha, A, &lda, &beta, C, &ldc); } -} // namespace LapackConnector -#endif // LAPACK_CONNECTOR_HPP + + static inline + double lange(const char& norm, const int& m, const int& n, const double* A, const int& lda, double* work) + { + return dlange_(&norm, &m, &n, A, &lda, work); + } + + static inline + double lange(const char& norm, const int& m, const int& n, const std::complex* A, const int& lda, double* work) + { + return zlange_(&norm, &m, &n, A, &lda, work); + } + +}; +#endif // LAPACKCONNECTOR_HPP diff --git a/source/source_base/module_external/scalapack_connector.h b/source/source_base/module_external/scalapack_connector.h index aa99babca0..a6e2ef5ea4 100644 --- a/source/source_base/module_external/scalapack_connector.h +++ b/source/source_base/module_external/scalapack_connector.h @@ -42,11 +42,23 @@ extern "C" void pdtran_(const int* m, const int* n, const double* alpha, const double* a, const int* ia, const int* ja, const int* desca, const double* beta, double* c, const int* ic, const int* jc, const int* descc); - void pztranu_(const int *m,const int*n, const std::complex* alpha, const std::complex* a, const int* ia, const int* ja, const int* desca, - const std::complex *beta , std::complex *c , const int *ic ,const int *jc ,const int *descc); - + const std::complex *beta , std::complex *c, const int *ic, const int *jc, const int *descc); + void pztranc_(const int *m, const int *n, + const std::complex *alpha, const std::complex *a, const int *ia, const int *ja, const int *desca, + const std::complex *beta, std::complex *c, const int *ic, const int *jc, const int *descc); + + double pdlange_(const char* norm, + const int* m, const int* n, + const double* a, const int* ia, const int* ja, const int* desca, + double* work); + + double pzlange_(const char* norm, + const int* m, const int* n, + const std::complex* a, const int* ia, const int* ja, const int* desca, + double* work); + void pzgemv_( const char *transa, const int *M, const int *N, @@ -144,13 +156,6 @@ extern "C" const std::complex *beta, const std::complex *c, const int *ic, const int *jc, const int *descc); - void pztranc_( - const int *M, const int *N, - const std::complex *alpha, - const std::complex *A, const int *IA, const int *JA, const int *DESCA, - const std::complex *beta, - std::complex *C, const int *IC, const int *JC, const int *DESCC); - void pdgemr2d_(const int *M, const int *N, double *A, const int *IA, const int *JA, const int *DESCA, double *B, const int *IB, const int *JB, const int *DESCB, @@ -394,11 +399,54 @@ class ScalapackConnector static inline void tranu( const int m, const int n, - const std::complex alpha , std::complex *a , const int ia , const int ja , const int *desca, - const std::complex beta , std::complex *c , const int ic , const int jc , const int *descc) + const double alpha, double *a, const int ia, const int ja, const int *desca, + const double beta, double *c, const int ic, const int jc, const int *descc) + { + pdtran_(&m, &n, &alpha, a, &ia, &ja, desca, &beta, c, &ic, &jc, descc); + } + static inline + void tranu( + const int m, const int n, + const std::complex alpha, std::complex *a, const int ia, const int ja, const int *desca, + const std::complex beta, std::complex *c, const int ic, const int jc, const int *descc) { pztranu_(&m, &n, &alpha, a, &ia, &ja, desca, &beta, c, &ic, &jc, descc); } + static inline + void tranc( + const int m, const int n, + const double alpha, double *a, const int ia, const int ja, const int *desca, + const double beta, double *c, const int ic, const int jc, const int *descc) + { + pdtran_(&m, &n, &alpha, a, &ia, &ja, desca, &beta, c, &ic, &jc, descc); + } + static inline + void tranc( + const int m, const int n, + const std::complex alpha, std::complex *a, const int ia, const int ja, const int *desca, + const std::complex beta, std::complex *c, const int ic, const int jc, const int *descc) + { + pztranc_(&m, &n, &alpha, a, &ia, &ja, desca, &beta, c, &ic, &jc, descc); + } + + static inline + double lange( + const char norm, + const int m, const int n, + const double* a, const int ia, const int ja, const int* desca, + double* work) + { + return pdlange_(&norm, &m, &n, a, &ia, &ja, desca, work); + } + static inline + double lange( + const char norm, + const int m, const int n, + const std::complex* a, const int ia, const int ja, const int* desca, + double* work) + { + return pzlange_(&norm, &m, &n, a, &ia, &ja, desca, work); + } static inline int potrf(char uplo, int na, double* U, int* desc) diff --git a/source/source_base/parallel_2d.cpp b/source/source_base/parallel_2d.cpp index b5bab96694..0c39ad9709 100644 --- a/source/source_base/parallel_2d.cpp +++ b/source/source_base/parallel_2d.cpp @@ -6,11 +6,23 @@ #include #include -bool Parallel_2D::in_this_processor(const int iw1_all, const int iw2_all) const +bool Parallel_2D::in_this_processor(const std::size_t iw1_all, const std::size_t iw2_all) const { return global2local_row(iw1_all) != -1 && global2local_col(iw2_all) != -1; } +int Parallel_2D::owner_processor(const std::size_t iw1_all, const std::size_t iw2_all) const +{ + assert(iw1_all < get_global_row_size() && iw2_all < get_global_col_size()); + if (is_serial) + { + return 0; + } + int proc_row = (iw1_all / nb) % dim0; + int proc_col = (iw2_all / nb) % dim1; + return proc_row * dim1 + proc_col; +} + int Parallel_2D::get_global_row_size() const { if (!is_serial) diff --git a/source/source_base/parallel_2d.h b/source/source_base/parallel_2d.h index 81fab1128e..2d89d12c4b 100644 --- a/source/source_base/parallel_2d.h +++ b/source/source_base/parallel_2d.h @@ -1,6 +1,7 @@ #ifndef _PARALLEL_2D_H_ #define _PARALLEL_2D_H_ +#include #include #include @@ -42,13 +43,13 @@ class Parallel_2D }; /// get the local index of a global index (row) - int global2local_row(const int igr) const + int global2local_row(const std::size_t igr) const { return global2local_row_[igr]; } /// get the local index of a global index (col) - int global2local_col(const int igc) const + int global2local_col(const std::size_t igc) const { return global2local_col_[igc]; } @@ -66,7 +67,10 @@ class Parallel_2D } /// check whether a global index is in this process - bool in_this_processor(const int iw1_all, const int iw2_all) const; + bool in_this_processor(const std::size_t iw1_all, const std::size_t iw2_all) const; + + /// get the owner processor of a global index + int owner_processor(const std::size_t iw1_all, const std::size_t iw2_all) const; /// side length of 2d square block int get_block_size() const diff --git a/source/source_io/module_ctrl/ctrl_scf_lcao.cpp b/source/source_io/module_ctrl/ctrl_scf_lcao.cpp index aafeb18e6a..08e4b0128d 100644 --- a/source/source_io/module_ctrl/ctrl_scf_lcao.cpp +++ b/source/source_io/module_ctrl/ctrl_scf_lcao.cpp @@ -637,6 +637,8 @@ void ModuleIO::ctrl_scf_lcao(UnitCell& ucell, { RPA_LRI rpa_lri_double(GlobalC::exx_info.info_ri); rpa_lri_double.postSCF(ucell, MPI_COMM_WORLD, *dm, pelec, kv, orb, pv, *psi); + if (inp.rpa_out_vel) + rpa_lri_double.out_velocity(ucell, gd, two_center_bundle, pv, *psi, pelec); } #endif diff --git a/source/source_io/module_parameter/input_parameter.h b/source/source_io/module_parameter/input_parameter.h index 62cc531409..be668e0bd2 100644 --- a/source/source_io/module_parameter/input_parameter.h +++ b/source/source_io/module_parameter/input_parameter.h @@ -367,11 +367,30 @@ struct Input_para std::vector abs_wavelen_range = {}; ///< the range of wavelength(nm) to output the absorption spectrum double abs_broadening = 0.01; ///< the broadening (eta) for LR-TDDFT absorption spectrum std::string abs_gauge - = "length"; ///< whether to use length or velocity gauge to calculate the absorption spectrum in LR-TDDFT + = "velocity"; ///< whether to use length or velocity gauge to calculate the absorption spectrum in LR-TDDFT std::string ri_hartree_benchmark = "none"; ///< whether to use the RI approximation for the Hartree potential in ///< LR-TDDFT for benchmark (with FHI-aims/ABACUS read-in style) std::vector aims_nbasis = {}; ///< the number of basis functions for each atom type used in FHI-aims (for benchmark) + std::string bse_tda = "tda"; ///< TDA type can be: "tda", "full", "both" + std::vector bse_spin_types = {"singlet", "triplet"}; ///< spin types for close-shell case to be calculated + bool bse_mem_save = false; ///< whether to save memory by adding V and W to BSE matrix directly + bool bse_ri_hartree = true; ///< whether to use RI approximation for Hartree term in BSE + int bse_use_fine_kgrid = 0; ///< 0: coarse k-grid; 1: uniform fine k-grid; 2: non-uniform fine k-grid + int bse_q_approx_mode = 0; ///< q→kpair mapping mode: 0=exact, 1=coarse q grid, 2=mixed + double bse_q_approx_threshold = 0.1; ///< threshold radius (Bohr^-1) for exact q in mode 2 + bool out_bse_ab = false; ///< whether to output the AB matrix to file + int bse_continue = 0; ///< which step to continue from previous BSE calculation + ///< 0: new; 1: continue from A_V; 2: A_V and A_W; 3: A_V, A_W and B_V; 4: A_V, A_W, B_V and B_W + int plot_istate = 0; ///< the index of excited state to be plotted (starting from 0) + std::string exciton_plot_type = "average"; ///< exciton density plot type: "average" or "conditional" + std::string exciton_plot_format = "cube"; ///< exciton plot format: "cube", "slice", or "both" + std::vector exciton_fixed_coordinate = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0}; ///< fixed hole and electron coordinates (Bohr): hx hy hz ex ey ez + std::string exciton_slice_plane = "ab"; ///< cross-section plane for conditional density: "ab", "bc", "ca" + double exciton_slice_pos = 0.0; ///< offset along perpendicular direction (Bohr) for slice + int exciton_slice_npoints = 200; ///< grid points per dimension for slice + double exciton_slice_scale = 1.3; ///< scale relative to BvK supercell for slice + // ============== #Parameters (11.Output) =========================== bool out_stru = false; ///< outut stru file each ion step int out_freq_elec = 0; ///< the frequency of electronic iter to output charge and wavefunction @@ -432,6 +451,8 @@ struct Input_para std::vector cal_symm_repr = {0, 3}; ///< output the symmetry representation matrix bool restart_save = false; ///< restart //Peize Lin add 2020-04-04 bool rpa = false; ///< rpa calculation + bool rpa_out_vel = false; ///< whether to output velocity matrix for librpa + std::string rpa_outdir = "./OUT.librpa/"; ///< output directory for librpa std::vector out_pchg = {}; ///< specify the bands to be calculated for partial charge std::vector out_wfc_norm = {}; ///< specify the bands to be calculated for norm of wfc std::vector out_wfc_re_im = {}; ///< specify the bands to be calculated for real and imaginary parts of wfc diff --git a/source/source_io/module_parameter/read_input_item_output.cpp b/source/source_io/module_parameter/read_input_item_output.cpp index 99591a4d27..f3dbd5196a 100644 --- a/source/source_io/module_parameter/read_input_item_output.cpp +++ b/source/source_io/module_parameter/read_input_item_output.cpp @@ -1330,10 +1330,37 @@ If EXX(exact exchange) is calculated (i.e. dft_fuctional==hse/hf/pbe0/scan0 or r "irreducible_sector.txt, symrot_k.txt and symrot_R.txt."; item.default_value = "False"; item.unit = ""; - item.availability = ""; + item.availability = "Numerical atomic orbital basis"; read_sync_bool(input.rpa); this->add_item(item); } + { + Input_Item item("rpa_out_vel"); + item.annotation = "whether to output velocity matrix for librpa"; + item.category = "Input files"; + item.type = "Boolean"; + item.description = " Velocity matrix in KS basis (in unit of eV *Angstrom). Loop layer: spin -> k -> direction -> KS_basis1 -> KS_basis2."; + item.default_value = "False"; + item.unit = "eV * A"; + item.availability = "Numerical atomic orbital basis"; + read_sync_bool(input.rpa_out_vel); + this->add_item(item); + } + { + Input_Item item("rpa_outdir"); + item.annotation = "output directory for librpa"; + item.category = "Input files"; + item.type = "String"; + item.description = "The directory to save files for LibRPA."; + item.default_value = "\"./OUT.librpa/\""; + item.unit = ""; + item.availability = "Numerical atomic orbital basis"; + read_sync_string(input.rpa_outdir); + item.reset_value = [](const Input_Item& item, Parameter& para) { + para.input.rpa_outdir = to_dir(para.input.rpa_outdir); + }; + this->add_item(item); + } { Input_Item item("out_pchg"); item.annotation = "specify the bands to be calculated for the partial (band-decomposed) charge densities"; diff --git a/source/source_io/module_parameter/read_input_item_tddft.cpp b/source/source_io/module_parameter/read_input_item_tddft.cpp index 18e8ad3c00..28650169a3 100644 --- a/source/source_io/module_parameter/read_input_item_tddft.cpp +++ b/source/source_io/module_parameter/read_input_item_tddft.cpp @@ -750,5 +750,129 @@ void ReadInput::item_lr_tddft() read_sync_double(input.abs_broadening); this->add_item(item); } + { + Input_Item item("bse_tda"); + item.annotation = "whether Tamm-Dancoff Approximation is used (can be 'tda', 'full' or 'both')"; + read_sync_string(input.bse_tda); + this->add_item(item); + } + { + Input_Item item("bse_spin_types"); + item.annotation = "which spin type is calculated (can be 'singlet', 'triplet', also for test 'rpa', 'ipa')"; + + item.read_value = [](const Input_Item& item, Parameter& para) { + size_t count = item.get_size(); + auto& ist = para.input.bse_spin_types; + ist.clear(); + for (int i = 0; i < count; i++) { ist.push_back(item.str_values[i]); } + }; + item.reset_value = [](const Input_Item& item, Parameter& para) { + if (para.input.bse_spin_types.empty()) { para.input.bse_spin_types={"singlet","triplet"}; } + }; + sync_stringvec(input.bse_spin_types, para.input.bse_spin_types.size(), "singlet"); + this->add_item(item); + } + { + Input_Item item("bse_mem_save"); + item.annotation = "whether to save memory by adding V and W to BSE matrix directly"; + item.reset_value = [](const Input_Item& item, Parameter& para) { + if (para.input.bse_mem_save == true) { para.input.bse_continue=0; para.input.bse_ri_hartree=true; } + }; + read_sync_bool(input.bse_mem_save); + this->add_item(item); + } + { + Input_Item item("bse_ri_hartree"); + item.annotation = "whether to use RI approximation for Hartree term in BSE"; + read_sync_bool(input.bse_ri_hartree); + this->add_item(item); + } + { + Input_Item item("bse_use_fine_kgrid"); + item.annotation = "whether to use a finer k-grid for BSE"; + read_sync_int(input.bse_use_fine_kgrid); + this->add_item(item); + } + { + Input_Item item("bse_q_approx_mode"); + item.annotation = "q->kpair mapping mode: 0=exact, 1=coarse q grid, 2=mixed"; + read_sync_int(input.bse_q_approx_mode); + this->add_item(item); + } + { + Input_Item item("bse_q_approx_threshold"); + item.annotation = "threshold radius (Bohr^-1) for exact q mapping in mode 2"; + read_sync_double(input.bse_q_approx_threshold); + this->add_item(item); + } + { + Input_Item item("out_bse_ab"); + item.annotation = "whether to output the AB matrix to file"; + read_sync_bool(input.out_bse_ab); + this->add_item(item); + } + { + Input_Item item("bse_continue"); + item.annotation = "which step to continue from previous BSE calculation"; + read_sync_int(input.bse_continue); + this->add_item(item); + } + { + Input_Item item("plot_istate"); + item.annotation = "which state of exciton to be ploted"; + read_sync_int(input.plot_istate); + this->add_item(item); + } + { + Input_Item item("exciton_plot_type"); + item.annotation = "exciton density plot type: average or conditional"; + read_sync_string(input.exciton_plot_type); + this->add_item(item); + } + { + Input_Item item("exciton_plot_format"); + item.annotation = "exciton plot format: cube, slice, or both"; + read_sync_string(input.exciton_plot_format); + this->add_item(item); + } + { + Input_Item item("exciton_fixed_coordinate"); + item.annotation = "fixed hole and electron Cartesian coordinates (Bohr): hx hy hz ex ey ez"; + item.read_value = [](const Input_Item& item, Parameter& para) { + parse_expression(item.str_values, para.input.exciton_fixed_coordinate); + }; + item.check_value = [](const Input_Item&, const Parameter& para) { + if (para.input.exciton_fixed_coordinate.size() != 6) + { + ModuleBase::WARNING_QUIT("ReadInput", "exciton_fixed_coordinate should have exactly 6 values"); + } + }; + sync_doublevec(input.exciton_fixed_coordinate, 6, 0.0); + this->add_item(item); + } + { + Input_Item item("exciton_slice_plane"); + item.annotation = "cross-section plane: ab, bc, or ca"; + read_sync_string(input.exciton_slice_plane); + this->add_item(item); + } + { + Input_Item item("exciton_slice_pos"); + item.annotation = "offset along perpendicular direction (Bohr) for slice"; + read_sync_double(input.exciton_slice_pos); + this->add_item(item); + } + { + Input_Item item("exciton_slice_npoints"); + item.annotation = "grid points per dimension for slice"; + read_sync_int(input.exciton_slice_npoints); + this->add_item(item); + } + { + Input_Item item("exciton_slice_scale"); + item.annotation = "scale relative to BvK supercell for slice"; + read_sync_double(input.exciton_slice_scale); + this->add_item(item); + } } } diff --git a/source/source_io/test/read_input_ptest.cpp b/source/source_io/test/read_input_ptest.cpp index 05f5b1e75f..ee3c28b24f 100644 --- a/source/source_io/test/read_input_ptest.cpp +++ b/source/source_io/test/read_input_ptest.cpp @@ -353,6 +353,8 @@ TEST_F(InputParaTest, ParaRead) EXPECT_EQ(param.inp.omc, 0); EXPECT_FALSE(param.inp.dft_plus_dmft); EXPECT_FALSE(param.inp.rpa); + EXPECT_FALSE(param.inp.rpa_out_vel); + EXPECT_EQ(param.inp.rpa_outdir, "./OUT.librpa/"); EXPECT_EQ(param.inp.imp_sol, 0); EXPECT_DOUBLE_EQ(param.inp.eb_k, 80.0); EXPECT_DOUBLE_EQ(param.inp.tau, 1.0798 * 1e-5); @@ -448,7 +450,27 @@ TEST_F(InputParaTest, ParaRead) EXPECT_EQ(param.inp.abs_wavelen_range.size(), 2); EXPECT_DOUBLE_EQ(param.inp.abs_wavelen_range[0], 0.0); EXPECT_DOUBLE_EQ(param.inp.abs_broadening, 0.01); - EXPECT_EQ(param.inp.abs_gauge, "length"); + EXPECT_EQ(param.inp.abs_gauge, "velocity"); + EXPECT_EQ(param.inp.bse_tda, "tda"); + EXPECT_EQ(param.inp.bse_spin_types[0], "singlet"); + EXPECT_EQ(param.inp.bse_spin_types[1], "triplet"); + EXPECT_FALSE(param.inp.bse_mem_save); + EXPECT_TRUE(param.inp.bse_ri_hartree); + EXPECT_EQ(param.inp.bse_use_fine_kgrid, 0); + EXPECT_FALSE(param.inp.out_bse_ab); + EXPECT_EQ(param.inp.bse_continue, 0); + EXPECT_EQ(param.inp.plot_istate, 0); + EXPECT_EQ(param.inp.exciton_plot_type, "average"); + EXPECT_EQ(param.inp.exciton_plot_format, "cube"); + ASSERT_EQ(param.inp.exciton_fixed_coordinate.size(), 6); + for (const double coordinate : param.inp.exciton_fixed_coordinate) + { + EXPECT_DOUBLE_EQ(coordinate, 0.0); + } + EXPECT_EQ(param.inp.exciton_slice_plane, "ab"); + EXPECT_DOUBLE_EQ(param.inp.exciton_slice_pos, 0.0); + EXPECT_EQ(param.inp.exciton_slice_npoints, 200); + EXPECT_DOUBLE_EQ(param.inp.exciton_slice_scale, 1.3); EXPECT_EQ(param.inp.rdmft, 0); EXPECT_DOUBLE_EQ(param.inp.rdmft_power_alpha, 0.656); } diff --git a/source/source_lcao/module_lr/CMakeLists.txt b/source/source_lcao/module_lr/CMakeLists.txt index 331e7dc22e..64f06765ac 100644 --- a/source/source_lcao/module_lr/CMakeLists.txt +++ b/source/source_lcao/module_lr/CMakeLists.txt @@ -4,9 +4,15 @@ if(ENABLE_LCAO) add_subdirectory(dm_trans) add_subdirectory(ri_benchmark) + if(ENABLE_LIBRI) + add_subdirectory(bse) + endif() + list(APPEND objects utils/lr_util.cpp utils/lr_util_hcontainer.cpp + utils/lr_io.cpp + utils/exciton_plotter.cpp ao_to_mo_transformer/ao_to_mo_parallel.cpp ao_to_mo_transformer/ao_to_mo_serial.cpp dm_trans/dm_trans_parallel.cpp @@ -20,6 +26,14 @@ if(ENABLE_LCAO) esolver_lrtd_lcao.cpp potentials/xc_kernel.cpp) + if(ENABLE_LIBRI) + list(APPEND objects + bse/esolver_bse_lcao.cpp + bse/hamilt_bse.cpp + bse/hamilt_bse_solver.cpp + bse/bse_util.cpp) + endif() + add_library( lr OBJECT diff --git a/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h b/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h index b6e87862a8..0cfcba9348 100644 --- a/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h +++ b/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h @@ -1,24 +1,22 @@ #pragma once #include #include "source_psi/psi.h" +#include "source_lcao/module_lr/utils/mo_type.h" #include +#include #ifdef __MPI #include "source_base/parallel_2d.h" #endif namespace LR { -#ifndef MO_TYPE_H -#define MO_TYPE_H - enum MO_TYPE { OO, VO, VV }; -#endif template - void ao_to_mo_forloop_serial( + void ao_to_mo_forloop_serial( const std::vector& mat_ao, const psi::Psi& coeff, const int& nocc, const int& nvirt, T* const mat_mo, - const MO_TYPE type = VO); + const LR_Util::MO_TYPE type = LR_Util::VO); template void ao_to_mo_blas( const std::vector& mat_ao, @@ -27,7 +25,7 @@ namespace LR const int& nvirt, T* const mat_mo, const bool add_on = true, - const MO_TYPE type = VO); + const LR_Util::MO_TYPE type = LR_Util::VO); #ifdef __MPI template void ao_to_mo_pblas( @@ -41,6 +39,6 @@ namespace LR const Parallel_2D& pmat_mo, T* const mat_mo, const bool add_on = true, - const MO_TYPE type = VO); + const LR_Util::MO_TYPE type = LR_Util::VO); #endif } \ No newline at end of file diff --git a/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_parallel.cpp b/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_parallel.cpp index 7d936365ce..f4d1bb79de 100644 --- a/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_parallel.cpp +++ b/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_parallel.cpp @@ -21,20 +21,21 @@ namespace LR const Parallel_2D& pmat_mo, double* mat_mo, const bool add_on, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { - ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_pblas"); + ModuleBase::TITLE("LR", "ao_to_mo_pblas"); assert(pmat_ao.comm() == pcoeff.comm() && pmat_ao.comm() == pmat_mo.comm()); assert(pmat_ao.blacs_ctxt == pcoeff.blacs_ctxt && pmat_ao.blacs_ctxt == pmat_mo.blacs_ctxt); assert(pmat_mo.get_local_size() > 0); const int nks = mat_ao.size(); + int nmo1_set, nmo2_set, imo1_set, imo2_set; + LR_Util::set_dim(type, nocc, nvirt, nmo1_set, nmo2_set, imo1_set, imo2_set); + const int nmo1 = nmo1_set; + const int nmo2 = nmo2_set; + const int imo1 = imo1_set + 1; + const int imo2 = imo2_set + 1; const int i1 = 1; - const int ivirt = nocc + 1; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? ivirt : i1; - const int imo2 = type == MO_TYPE::OO ? i1 : ivirt; Parallel_2D pVc; // for intermediate Vc LR_Util::setup_2d_division(pVc, pmat_ao.get_block_size(), naos, nmo1, pmat_ao.blacs_ctxt); @@ -79,20 +80,21 @@ namespace LR const Parallel_2D& pmat_mo, std::complex* const mat_mo, const bool add_on, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { - ModuleBase::TITLE("hamilt_lrtd", "cal_AX_plas"); + ModuleBase::TITLE("LR", "ao_to_mo_pblas"); assert(pmat_ao.comm() == pcoeff.comm() && pmat_ao.comm() == pmat_mo.comm()); assert(pmat_ao.blacs_ctxt == pcoeff.blacs_ctxt && pmat_ao.blacs_ctxt == pmat_mo.blacs_ctxt); assert(pmat_mo.get_local_size() > 0); const int nks = mat_ao.size(); + int nmo1_set, nmo2_set, imo1_set, imo2_set; + LR_Util::set_dim(type, nocc, nvirt, nmo1_set, nmo2_set, imo1_set, imo2_set); + const int nmo1 = nmo1_set; + const int nmo2 = nmo2_set; + const int imo1 = imo1_set + 1; + const int imo2 = imo2_set + 1; const int i1 = 1; - const int ivirt = nocc + 1; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? ivirt : i1; - const int imo2 = type == MO_TYPE::OO ? i1 : ivirt; Parallel_2D pVc; // for intermediate Vc LR_Util::setup_2d_division(pVc, pmat_ao.get_block_size(), naos, nmo1, pmat_ao.blacs_ctxt); diff --git a/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_serial.cpp b/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_serial.cpp index b65f871f32..74afc6e803 100644 --- a/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_serial.cpp +++ b/source/source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo_serial.cpp @@ -11,15 +11,13 @@ namespace LR const int& nocc, const int& nvirt, double* mat_mo, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { - ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_forloop_serial"); + ModuleBase::TITLE("LR", "ao_to_mo_forloop_serial"); const int nks = mat_ao.size(); const int naos = coeff.get_nbasis(); - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); ModuleBase::GlobalFunc::ZEROS(mat_mo, nks * nmo1 * nmo2); @@ -49,15 +47,13 @@ namespace LR const int& nocc, const int& nvirt, std::complex* const mat_mo, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { - ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_forloop_serial"); + ModuleBase::TITLE("LR", "ao_to_mo_forloop_serial"); const int nks = mat_ao.size(); const int naos = coeff.get_nbasis(); - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); ModuleBase::GlobalFunc::ZEROS(mat_mo, nks * nmo1 * nmo2); @@ -88,15 +84,13 @@ namespace LR const int& nvirt, double* mat_mo, const bool add_on, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { - ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_blas"); + ModuleBase::TITLE("LR", "ao_to_mo_blas"); const int nks = mat_ao.size(); - const int naos = coeff.get_nbasis(); - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; + const int naos = coeff.get_nbasis(); + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); for (int isk = 0;isk < nks;++isk) { @@ -129,15 +123,13 @@ namespace LR const int& nvirt, std::complex* const mat_mo, const bool add_on, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { - ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_blas"); + ModuleBase::TITLE("LR", "ao_to_mo_blas"); const int nks = mat_ao.size(); const int naos = coeff.get_nbasis(); - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); for (int isk = 0;isk < nks;++isk) { diff --git a/source/source_lcao/module_lr/ao_to_mo_transformer/test/ao_to_mo_test.cpp b/source/source_lcao/module_lr/ao_to_mo_transformer/test/ao_to_mo_test.cpp index b8f5bbefc7..09048eed31 100644 --- a/source/source_lcao/module_lr/ao_to_mo_transformer/test/ao_to_mo_test.cpp +++ b/source/source_lcao/module_lr/ao_to_mo_transformer/test/ao_to_mo_test.cpp @@ -81,10 +81,10 @@ TEST_F(AO2MOTest, DoubleSerial) for (auto& v : V) { set_rand(v.data(), size_v); } LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &vo_for(istate, 0, 0)); LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &vo_blas(istate, 0, 0), false); - LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &oo_for(istate, 0, 0), LR::MO_TYPE::OO); - LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &oo_blas(istate, 0, 0), false, LR::MO_TYPE::OO); - LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &vv_for(istate, 0, 0), LR::MO_TYPE::VV); - LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &vv_blas(istate, 0, 0), false, LR::MO_TYPE::VV); + LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &oo_for(istate, 0, 0), LR_Util::MO_TYPE::OO); + LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &oo_blas(istate, 0, 0), false, LR_Util::MO_TYPE::OO); + LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &vv_for(istate, 0, 0), LR_Util::MO_TYPE::VV); + LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &vv_blas(istate, 0, 0), false, LR_Util::MO_TYPE::VV); } check_eq(&vo_for(0, 0, 0), &vo_blas(0, 0, 0), nstate * s.nks * s.nocc * s.nvirt); check_eq(&oo_for(0, 0, 0), &oo_blas(0, 0, 0), nstate * s.nks * s.nocc * s.nocc); @@ -113,10 +113,10 @@ TEST_F(AO2MOTest, ComplexSerial) for (auto& v : V) { set_rand(v.data>(), size_v); } LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &vo_for(istate, 0, 0)); LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &vo_blas(istate, 0, 0), false); - LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &oo_for(istate, 0, 0), LR::MO_TYPE::OO); - LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &oo_blas(istate, 0, 0), false, LR::MO_TYPE::OO); - LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &vv_for(istate, 0, 0), LR::MO_TYPE::VV); - LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &vv_blas(istate, 0, 0), false, LR::MO_TYPE::VV); + LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &oo_for(istate, 0, 0), LR_Util::MO_TYPE::OO); + LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &oo_blas(istate, 0, 0), false, LR_Util::MO_TYPE::OO); + LR::ao_to_mo_forloop_serial(V, c, s.nocc, s.nvirt, &vv_for(istate, 0, 0), LR_Util::MO_TYPE::VV); + LR::ao_to_mo_blas(V, c, s.nocc, s.nvirt, &vv_blas(istate, 0, 0), false, LR_Util::MO_TYPE::VV); } check_eq(&vo_for(0, 0, 0), &vo_blas(0, 0, 0), nstate * s.nks * s.nocc * s.nvirt); check_eq(&oo_for(0, 0, 0), &oo_blas(0, 0, 0), nstate * s.nks * s.nocc * s.nocc); @@ -150,10 +150,13 @@ TEST_F(AO2MOTest, DoubleParallel) EXPECT_GE(s.naos, pc.dim0); psi::Psi vo_pblas_loc(s.nks, nstate, pvo.get_local_size(), pvo.get_local_size(), false); psi::Psi vo_gather(s.nks, nstate, s.nocc * s.nvirt, s.nocc * s.nvirt, false); + vo_gather.zero_out(); psi::Psi oo_pblas_loc(s.nks, nstate, poo.get_local_size(), poo.get_local_size(), false); psi::Psi oo_gather(s.nks, nstate, s.nocc * s.nocc, s.nocc * s.nocc, false); + oo_gather.zero_out(); psi::Psi vv_pblas_loc(s.nks, nstate, pvv.get_local_size(), pvv.get_local_size(), false); psi::Psi vv_gather(s.nks, nstate, s.nvirt * s.nvirt, s.nvirt * s.nvirt, false); + vv_gather.zero_out(); for (int istate = 0;istate < nstate;++istate) { for (int isk = 0;isk < s.nks;++isk) @@ -162,21 +165,23 @@ TEST_F(AO2MOTest, DoubleParallel) set_rand(&c(isk, 0, 0), pc.get_local_size()); } LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, pvo, &vo_pblas_loc(istate, 0, 0), false); - LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, poo, &oo_pblas_loc(istate, 0, 0), false, LR::MO_TYPE::OO); - LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, pvv, &vv_pblas_loc(istate, 0, 0), false, LR::MO_TYPE::VV); + LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, poo, &oo_pblas_loc(istate, 0, 0), false, LR_Util::MO_TYPE::OO); + LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, pvv, &vv_pblas_loc(istate, 0, 0), false, LR_Util::MO_TYPE::VV); // gather AX and output for (int isk = 0;isk < s.nks;++isk) { - LR_Util::gather_2d_to_full(pvo, &vo_pblas_loc(istate, isk, 0), &vo_gather(istate, isk, 0), false/*pblas: row first*/, s.nvirt, s.nocc); - LR_Util::gather_2d_to_full(poo, &oo_pblas_loc(istate, isk, 0), &oo_gather(istate, isk, 0), false/*pblas: row first*/, s.nocc, s.nocc); - LR_Util::gather_2d_to_full(pvv, &vv_pblas_loc(istate, isk, 0), &vv_gather(istate, isk, 0), false/*pblas: row first*/, s.nvirt, s.nvirt); + LR_Util::gather_2d_to_full(pvo, &vo_pblas_loc(istate, isk, 0), &vo_gather(istate, isk, 0), false/*pblas: col major*/, s.nvirt, s.nocc); + LR_Util::gather_2d_to_full(poo, &oo_pblas_loc(istate, isk, 0), &oo_gather(istate, isk, 0), false/*pblas: col major*/, s.nocc, s.nocc); + LR_Util::gather_2d_to_full(pvv, &vv_pblas_loc(istate, isk, 0), &vv_gather(istate, isk, 0), false/*pblas: col major*/, s.nvirt, s.nvirt); } // compare to global AX std::vector V_full(s.nks, container::Tensor(DAT::DT_DOUBLE, DEV::CpuDevice, { s.naos, s.naos })); std::vector ngk_temp_1(s.nks, s.naos); psi::Psi c_full(s.nks, s.nocc + s.nvirt, s.naos, ngk_temp_1, true); + c_full.zero_out(); for (int isk = 0;isk < s.nks;++isk) { + V_full[isk].zero(); LR_Util::gather_2d_to_full(pV, V.at(isk).data(), V_full.at(isk).data(), false, s.naos, s.naos); LR_Util::gather_2d_to_full(pc, &c(isk, 0, 0), &c_full(isk, 0, 0), false, s.naos, s.nocc + s.nvirt); } @@ -186,10 +191,10 @@ TEST_F(AO2MOTest, DoubleParallel) LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &vo_full_istate(0, 0, 0), false); check_eq(&vo_full_istate(0, 0, 0), &vo_gather(istate, 0, 0), s.nks * s.nocc * s.nvirt); psi::Psi oo_full_istate(s.nks, 1, s.nocc * s.nocc, s.nocc * s.nocc, false); - LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &oo_full_istate(0, 0, 0), false, LR::MO_TYPE::OO); + LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &oo_full_istate(0, 0, 0), false, LR_Util::MO_TYPE::OO); check_eq(&oo_full_istate(0, 0, 0), &oo_gather(istate, 0, 0), s.nks * s.nocc * s.nocc); psi::Psi vv_full_istate(s.nks, 1, s.nvirt * s.nvirt, s.nvirt * s.nvirt, false); - LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &vv_full_istate(0, 0, 0), false, LR::MO_TYPE::VV); + LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &vv_full_istate(0, 0, 0), false, LR_Util::MO_TYPE::VV); check_eq(&vv_full_istate(0, 0, 0), &vv_gather(istate, 0, 0), s.nks * s.nvirt * s.nvirt); } } @@ -216,10 +221,13 @@ TEST_F(AO2MOTest, ComplexParallel) psi::Psi> vo_pblas_loc(s.nks, nstate, pvo.get_local_size(), pvo.get_local_size(), false); psi::Psi> vo_gather(s.nks, nstate, s.nocc * s.nvirt, s.nocc * s.nvirt, false); + vo_gather.zero_out(); psi::Psi> oo_pblas_loc(s.nks, nstate, poo.get_local_size(), poo.get_local_size(), false); psi::Psi> oo_gather(s.nks, nstate, s.nocc * s.nocc, s.nocc * s.nocc, false); + oo_gather.zero_out(); psi::Psi> vv_pblas_loc(s.nks, nstate, pvv.get_local_size(), pvv.get_local_size(), false); psi::Psi> vv_gather(s.nks, nstate, s.nvirt * s.nvirt, s.nvirt * s.nvirt, false); + vv_gather.zero_out(); for (int istate = 0;istate < nstate;++istate) { for (int isk = 0;isk < s.nks;++isk) @@ -228,22 +236,24 @@ TEST_F(AO2MOTest, ComplexParallel) set_rand(&c(isk, 0, 0), pc.get_local_size()); } LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, pvo, &vo_pblas_loc(istate, 0, 0), false); - LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, poo, &oo_pblas_loc(istate, 0, 0), false, LR::MO_TYPE::OO); - LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, pvv, &vv_pblas_loc(istate, 0, 0), false, LR::MO_TYPE::VV); + LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, poo, &oo_pblas_loc(istate, 0, 0), false, LR_Util::MO_TYPE::OO); + LR::ao_to_mo_pblas(V, pV, c, pc, s.naos, s.nocc, s.nvirt, pvv, &vv_pblas_loc(istate, 0, 0), false, LR_Util::MO_TYPE::VV); // gather AX and output for (int isk = 0;isk < s.nks;++isk) { - LR_Util::gather_2d_to_full(pvo, &vo_pblas_loc(istate, isk, 0), &vo_gather(istate, isk, 0), false/*pblas: row first*/, s.nvirt, s.nocc); - LR_Util::gather_2d_to_full(poo, &oo_pblas_loc(istate, isk, 0), &oo_gather(istate, isk, 0), false/*pblas: row first*/, s.nocc, s.nocc); - LR_Util::gather_2d_to_full(pvv, &vv_pblas_loc(istate, isk, 0), &vv_gather(istate, isk, 0), false/*pblas: row first*/, s.nvirt, s.nvirt); + LR_Util::gather_2d_to_full(pvo, &vo_pblas_loc(istate, isk, 0), &vo_gather(istate, isk, 0), false/*pblas: col major*/, s.nvirt, s.nocc); + LR_Util::gather_2d_to_full(poo, &oo_pblas_loc(istate, isk, 0), &oo_gather(istate, isk, 0), false/*pblas: col major*/, s.nocc, s.nocc); + LR_Util::gather_2d_to_full(pvv, &vv_pblas_loc(istate, isk, 0), &vv_gather(istate, isk, 0), false/*pblas: col major*/, s.nvirt, s.nvirt); } // compare to global AX std::vector V_full(s.nks, container::Tensor(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, { s.naos, s.naos })); std::vector ngk_temp_2(s.nks, s.naos); psi::Psi> c_full(s.nks, s.nocc + s.nvirt, s.naos, ngk_temp_2, true); + c_full.zero_out(); for (int isk = 0;isk < s.nks;++isk) { + V_full[isk].zero(); LR_Util::gather_2d_to_full(pV, V.at(isk).data>(), V_full.at(isk).data>(), false, s.naos, s.naos); LR_Util::gather_2d_to_full(pc, &c(isk, 0, 0), &c_full(isk, 0, 0), false, s.naos, s.nocc + s.nvirt); } @@ -253,10 +263,10 @@ TEST_F(AO2MOTest, ComplexParallel) LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &vo_full_istate(0, 0, 0), false); check_eq(&vo_full_istate(0, 0, 0), &vo_gather(istate, 0, 0), s.nks * s.nocc * s.nvirt); psi::Psi> oo_full_istate(s.nks, 1, s.nocc * s.nocc, s.nocc * s.nvirt, false); - LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nocc, &oo_full_istate(0, 0, 0), false, LR::MO_TYPE::OO); + LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nocc, &oo_full_istate(0, 0, 0), false, LR_Util::MO_TYPE::OO); check_eq(&oo_full_istate(0, 0, 0), &oo_gather(istate, 0, 0), s.nks * s.nocc * s.nocc); psi::Psi> vv_full_istate(s.nks, 1, s.nvirt * s.nvirt, s.nocc * s.nvirt, false); - LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &vv_full_istate(0, 0, 0), false, LR::MO_TYPE::VV); + LR::ao_to_mo_blas(V_full, c_full, s.nocc, s.nvirt, &vv_full_istate(0, 0, 0), false, LR_Util::MO_TYPE::VV); check_eq(&vv_full_istate(0, 0, 0), &vv_gather(istate, 0, 0), s.nks * s.nvirt * s.nvirt); } } diff --git a/source/source_lcao/module_lr/bse/CMakeLists.txt b/source/source_lcao/module_lr/bse/CMakeLists.txt new file mode 100644 index 0000000000..b7bc0d0f89 --- /dev/null +++ b/source/source_lcao/module_lr/bse/CMakeLists.txt @@ -0,0 +1,5 @@ +if(ENABLE_LCAO AND ENABLE_LIBRI) + if(BUILD_TESTING) + add_subdirectory(test) + endif() +endif() diff --git a/source/source_lcao/module_lr/bse/bse_util.cpp b/source/source_lcao/module_lr/bse/bse_util.cpp new file mode 100644 index 0000000000..de1420f1b2 --- /dev/null +++ b/source/source_lcao/module_lr/bse/bse_util.cpp @@ -0,0 +1,151 @@ +#include "bse_util.h" +#include "source_lcao/module_lr/utils/lr_util.hpp" + +namespace BSE_Util +{ +#ifdef __MPI +template <> +container::Tensor cal_dm_trans_onebase_pblas( + const psi::Psi& c, + const Parallel_2D& pc, + const int& ik, + const int& naos, + const int& imo1, // imo1 → imo2, for excitation imo1[0,nocc), imo2[nocc,nocc+nvirt) + const int& imo2, + const Parallel_Orbitals& pmat, + const double& factor) +{ + ModuleBase::TITLE("BSE_Util", "cal_dm_trans_onebase_pblas(double)"); + assert(pc.comm() == pmat.comm()); + assert(pc.blacs_ctxt == pmat.blacs_ctxt); + assert(pmat.get_local_size() > 0); + assert(pmat.get_global_row_size() == naos); + assert(pmat.get_global_col_size() == naos); + c.fix_k(ik); + const int one = 1; + + container::Tensor dm_trans(DAT::DT_DOUBLE, DEV::CpuDevice, + { pmat.get_col_size(), pmat.get_row_size() }); // row is "inside"(memory contiguity) for pblas + + char transa = 'N', transb = 'T'; + const double beta = 0; + int imo1_ = imo1 + 1; // fortran index + int imo2_ = imo2 + 1; + // for excitation => [C_virt * C_occ^T]^T = C_occ * C_virt^T + // (lhs for row-major Tensor, rhs for column-major pblas, occ_aos is contiguous) + + ScalapackConnector::gemm(transa, transb, naos, naos, one, + factor, c.get_pointer(), one, imo1_, pc.desc, + c.get_pointer(), one, imo2_, pc.desc, + beta, dm_trans.data(), one, one, pmat.desc); + + return dm_trans; +} + +template <> +container::Tensor cal_dm_trans_onebase_pblas( + const psi::Psi>& c, + const Parallel_2D& pc, + const int& ik, + const int& naos, + const int& imo1, // imo1 → imo2, for excitation imo1[0,nocc), imo2[nocc,nocc+nvirt) + const int& imo2, + const Parallel_Orbitals& pmat, + const std::complex& factor) +{ + ModuleBase::TITLE("BSE_Util", "cal_dm_trans_onebase_pblas(complex)"); + assert(pc.comm() == pmat.comm()); + assert(pc.blacs_ctxt == pmat.blacs_ctxt); + assert(pmat.get_local_size() > 0); + assert(pmat.get_global_row_size() == naos); + assert(pmat.get_global_col_size() == naos); + c.fix_k(ik); + const int one = 1; + + container::Tensor dm_trans(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, + { pmat.get_col_size(), pmat.get_row_size() }); // row is "inside"(memory contiguity) for pblas + std::vector> dm_trans_tmp(pmat.get_local_size()); + char transa = 'N', transb = 'C'; + const std::complex alpha(1.0, 0.0); + const std::complex beta(0.0, 0.0); + int imo1_ = imo1 + 1; // fortran index + int imo2_ = imo2 + 1; + // for excitation => [C_virt * C_occ^\dagger]^T = C_occ^* * C_virt^T + // (lhs for row-major Tensor, rhs for column-major pblas, occ_aos is contiguous) + + // Since psi::get_pointer(ib) is only for global wfc, we implement the conj through pzgemm + // As a comparasion, see cal_dm_trans_onebase_blas(complex) + ScalapackConnector::gemm(transa, transb, naos, naos, one, + factor, c.get_pointer(), one, imo2_, pc.desc, + c.get_pointer(), one, imo1_, pc.desc, + beta, dm_trans_tmp.data(), one, one, pmat.desc); + + ScalapackConnector::tranu(naos, naos, alpha, + dm_trans_tmp.data(), one, one, pmat.desc, + beta, + dm_trans.data>(), one, one, pmat.desc); + + return dm_trans; +} +#endif + +template <> +container::Tensor cal_dm_trans_onebase_blas( + const psi::Psi& c, + const int& ik, + const int& naos, + const int& imo1, // imo1 → imo2, for excitation imo1[0,nocc), imo2[nocc,nocc+nvirt) + const int& imo2, + const double& factor) +{ + ModuleBase::TITLE("BSE_Util", "cal_dm_trans_onebase_blas(double)"); + c.fix_k(ik); + const int one = 1; + + container::Tensor dm_trans(DAT::DT_DOUBLE, DEV::CpuDevice, { naos, naos }); + + char transa = 'N', transb = 'T'; + const double beta = 0; + // for excitation => [C_virt * C_occ^T]^T = C_occ * C_virt^T + // (lhs for row-major Tensor, rhs for column-major pblas, occ_aos is contiguous) + + BlasConnector::gemm_cm(transa, transb, naos, naos, one, + factor, c.get_pointer(imo1), naos, + c.get_pointer(imo2), naos, + beta, dm_trans.data(), naos); + + return dm_trans; +} + +template <> +container::Tensor cal_dm_trans_onebase_blas( + const psi::Psi>& c, + const int& ik, + const int& naos, + const int& imo1, // imo1 → imo2, for excitation imo1[0,nocc), imo2[nocc,nocc+nvirt) + const int& imo2, + const std::complex& factor) +{ + ModuleBase::TITLE("BSE_Util", "cal_dm_trans_onebase_blas(complex)"); + c.fix_k(ik); + const int one = 1; + + container::Tensor dm_trans(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, { naos, naos }); + + char transa = 'N', transb = 'T'; + const std::complex beta = 0; + // for excitation => [C_virt * C_occ^\dagger]^T = C_occ^* * C_virt^T + // (lhs for row-major Tensor, rhs for column-major pblas, occ_aos is contiguous) + + std::vector> c_imo1_conj(naos); + for (int ibasis = 0; ibasis < naos; ++ibasis) { + c_imo1_conj[ibasis] = LR_Util::get_conj(c(imo1, ibasis)); + } + BlasConnector::gemm_cm(transa, transb, naos, naos, one, + factor, c_imo1_conj.data(), naos, + c.get_pointer(imo2), naos, + beta, dm_trans.data>(), naos); + + return dm_trans; +} +} \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/bse_util.h b/source/source_lcao/module_lr/bse/bse_util.h new file mode 100644 index 0000000000..ebf13e1604 --- /dev/null +++ b/source/source_lcao/module_lr/bse/bse_util.h @@ -0,0 +1,96 @@ +#pragma once + +#include +#include "source_psi/psi.h" +#include "source_base/tool_title.h" +#include "source_base/global_variable.h" +#include "source_base/module_external/blas_connector.h" +#include "source_base/module_external/scalapack_connector.h" +#include "source_base/parallel_2d.h" +#include "source_basis/module_ao/parallel_orbitals.h" +#include +#include +#include +namespace BSE_Util +{ +/// ================ info ================== +inline void print_mem_estimate(const std::string& name, + const size_t& local_size, + const size_t& type_size) +{ + double mem_MB = static_cast(local_size * type_size) / (1024.0 * 1024.0); + double mem_GB = mem_MB / 1024.0; + + GlobalV::ofs_running << "Allocating " << name << ", memory size: "; + if (mem_GB >= 1.0) { + GlobalV::ofs_running << std::fixed << std::setprecision(5) << mem_GB << " GB"; + } else { + GlobalV::ofs_running << std::fixed << std::setprecision(5) << mem_MB << " MB"; + } + GlobalV::ofs_running << std::endl; +} + +/// ================ RI ================== +using TA = int; +using TC = std::array; +using TAC = std::pair; +template +using TLRI = std::map>>; +template +bool move_R_tensor(TLRI& tensor_map, const TA ia, const TA ja, const TC& R_original, const TC& R_new) +{ + auto it_a = tensor_map.find(ia); + if (it_a == tensor_map.end()) return false; + + auto& map_ja = it_a->second; + + const TAC old_key{ja, R_original}; + auto it_b = map_ja.find(old_key); + if (it_b == map_ja.end()) return false; + + const TAC new_key{ja, R_new}; + assert(map_ja.count(new_key) == 0 && "Error: target R tensor already exists in move_R_tensor!"); + + map_ja[new_key] = std::move(it_b->second); + map_ja.erase(it_b); + return true; +} + +/// ================ Container =============== +using DAT = container::DataType; +using DEV = container::DeviceType; + +/// @brief Struct to get DataType enum for different tensor data types, not used currently +template +using DAT_ENUM = container::DataTypeToEnum; + +/// =============== Algorithm =================== + +/// ================ DM_onebase =================== +#ifdef __MPI +/// @brief calculate the 2d-block transition density matrix in AO basis +/// \f[ \tilde{\rho}_{\mu\mu}=c_{j,\mu}c^*_{b,\nu} \f] +template +container::Tensor cal_dm_trans_onebase_pblas( + const psi::Psi& c, + const Parallel_2D& pc, + const int& ik, + const int& naos, + const int& imo1, // imo1 → imo2, for excitation imo1[1,nocc], imo2[nocc+1,nocc+nvirt] + const int& imo2, + const Parallel_Orbitals& pmat, + const T& factor = (T)1.0); + +#endif + +/// @brief calculate the transition density matrix in AO basis +template +container::Tensor cal_dm_trans_onebase_blas( + const psi::Psi& c, + const int& ik, + const int& naos, + const int& imo1, // imo1 → imo2, for excitation imo1[1,nocc], imo2[nocc+1,nocc+nvirt] + const int& imo2, + const T& factor = (T)1.0); + +} diff --git a/source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp b/source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp new file mode 100644 index 0000000000..c83bbb2e67 --- /dev/null +++ b/source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp @@ -0,0 +1,731 @@ +#include "source_lcao/module_lr/bse/esolver_bse_lcao.h" +#include +#include "source_cell/module_neighbor/sltk_atom_arrange.h" +#include "source_io/module_output/print_info.h" +#include "source_lcao/module_gint/gint.h" +namespace BSE +{ +template +ESolver_BSE::ESolver_BSE(const Input_para& inp, UnitCell& ucell) : + LR::ESolver_LR(inp, ucell, LR::DeferredESolverLRInit{}) +{ + ModuleBase::TITLE("ESolver_BSE", "ESolver_BSE(from scratch)"); + ModuleBase::timer::start("ESolver_BSE", "constructor"); + // xc kernel + this->xc_kernel = LR_Util::tolower(inp.xc_kernel); + + // necessary steps in ESolver_FP + ModuleESolver::ESolver_FP::before_all_runners(ucell, inp); + this->pelec = new elecstate::ElecStateLCAO(); + + this->kRlist = LR_IO::RI_kRlist(this->ucell, &this->kv); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "Set K-POINTS and R-list for RI"); + ModuleIO::print_parameters(ucell, this->kv, inp); + + this->parameter_check(); + + /// read orbitals and build the interpolation table + this->two_center_bundle_.build_orb(ucell.ntype, ucell.orbital_fn.data(), inp.orbital_dir); + + this->two_center_bundle_.to_LCAO_Orbitals(this->orb_, inp.lcao_ecut, inp.lcao_dk, inp.lcao_dr, inp.lcao_rmax, + inp.out_element_info, inp.cal_force); + this->orb_cutoff_ = this->orb_.cutoffs(); + if (LR_Util::tolower(this->input.abs_gauge) == "velocity") + { + this->setup_2center_table(this->two_center_bundle_, this->orb_, ucell); + } + + this->set_dimension(); + // setup 2d-block distribution for AO-matrix and KS wfc + LR_Util::setup_2d_division(this->paraMat_, 1, this->nbasis, this->nbasis); +#ifdef __MPI + this->paraMat_.set_desc_wfc_Eij(this->nbasis, this->nbands, this->paraMat_.get_row_size()); + int err = this->paraMat_.set_nloc_wfc_Eij(this->nbands, GlobalV::ofs_running, GlobalV::ofs_warning); + this->paraMat_.set_atomic_trace(ucell.get_iat2iwt(), ucell.nat, this->nbasis); +#else + this->paraMat_.nrow_bands = this->nbasis; + this->paraMat_.ncol_bands = this->nbands; +#endif + + this->psi_ks = new psi::Psi(this->kv.get_nks(), + this->paraMat_.ncol_bands, + this->paraMat_.get_row_size(), + this->kv.ngk, + true); + this->psi_ks_global = new psi::Psi(this->kv.get_nks(), + this->nbands, + this->nbasis, + this->kv.ngk, + true); + this->read_ks_wfc(); + // NOTE: openshell is not implemented in BSE + if (this->nspin == 2) + { + this->nupdown = this->cal_nupdown_form_occ(this->pelec->wg); + this->reset_dim_spin2(); + } + + LR_Util::setup_2d_division(this->paraC_, this->paraMat_.get_block_size(), this->nbasis, this->nbands +#ifdef __MPI + , this->paraMat_.blacs_ctxt +#endif + ); + + this->Pgrid.init(this->pw_rho->nx, + this->pw_rho->ny, + this->pw_rho->nz, + this->pw_rho->nplane, + this->pw_rho->nrxx, + this->pw_big->nbz, + this->pw_big->bz); + + // search adjacent atoms and init Gint + double search_radius = -1.0; + search_radius = atom_arrange::set_sr_NL(GlobalV::ofs_running, + PARAM.inp.out_level, + this->orb_.get_rcutmax_Phi(), + ucell.infoNL->get_rcutmax_Beta(), + PARAM.globalv.gamma_only_local); + atom_arrange::search(PARAM.globalv.search_pbc, + GlobalV::ofs_running, + this->gd, + this->ucell, + search_radius, + PARAM.inp.test_atom_input); + + this->gint_info_.reset(new ModuleGint::GintInfo( + this->pw_big->nbx, + this->pw_big->nby, + this->pw_big->nbz, + this->pw_rho->nx, + this->pw_rho->ny, + this->pw_rho->nz, + 0, + 0, + this->pw_big->nbzp_start, + this->pw_big->nbx, + this->pw_big->nby, + this->pw_big->nbzp, + this->orb_.Phi, + ucell, + this->gd)); + ModuleGint::Gint::set_gint_info(this->gint_info_.get()); + + this->pot.resize(this->nspin, nullptr); + if (this->input.lr_solver != "spectrum" && this->input.lr_solver != "plot") + { + this->mo_lri = LR_Util::make_unique>(this->ucell, + this->nk, + this->kRlist, + this->nocc[0], + this->nvirt[0], + *this->psi_ks_global); + + if (!this->input.bse_ri_hartree && this->input.ri_hartree_benchmark == "none") + { + Charge chg_gs; + this->read_ks_chg(chg_gs); + this->init_pot(chg_gs); + } + } + ModuleBase::timer::end("ESolver_BSE", "constructor"); +} + +template +void ESolver_BSE::runner(UnitCell& ucell, const int istep) +{ + ModuleBase::TITLE("ESolver_BSE", "runner"); + ModuleBase::timer::start("ESolver_BSE", "runner"); + //allocate 2-particle state and setup 2d division + this->allocate_eigen_infos(); + + auto efile_out = [&](const std::string& label)->std::string { + return PARAM.globalv.global_out_dir + "Excitation_Energy_" + label + ".dat";}; + auto vfile_out = [&](const std::string& label)->std::string { + return PARAM.globalv.global_out_dir + "Excitation_Amplitude_" + label + "_" + std::to_string(GlobalV::MY_RANK) + ".dat";}; + auto efile_in = [&](const std::string& label)->std::string { + return PARAM.globalv.global_readin_dir + "Excitation_Energy_" + label + ".dat";}; + auto vfile_in = [&](const std::string& label)->std::string { + return PARAM.globalv.global_readin_dir + "Excitation_Amplitude_" + label + "_" + std::to_string(GlobalV::MY_RANK) + ".dat";}; + + if (this->input.lr_solver == "elpa") + { + if (this->input.bse_spin_types == std::vector{"ipa"}) + { + this->ipa_solver(); + } + else + { + std::cout << "Calculating Casida/BSE matrix directly." << std::endl; + assert(this->xc_kernel == "bse"); + this->lri_init(); + HamiltBSE bse_matrix(this->nspin, this->nbasis, this->nocc, this->nvirt, this->ucell, + this->orb_cutoff_, this->gd, *this->psi_ks, *this->psi_ks_global, this->eig_gw, + *this->mo_lri, + this->pot[0], this->kv, this->paraX_, this->paraC_, this->paraMat_, + this->input.bse_spin_types, + this->input.bse_tda, + this->input.ri_hartree_benchmark); + + auto write_tda_states = [&](const std::string& label, + const Real* e, + const T* v, + const int& dim, + const int& nst, + const int& prec = 8) -> void { + if (GlobalV::MY_RANK == 0) + { + assert(nst == LR_Util::write_value(efile_out(label), prec, e, nst)); + } + assert(nst * dim == LR_Util::write_value(vfile_out(label), prec, v, nst, dim)); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "write tda states " + label); + }; + auto write_full_states = [&](const std::string& label, + const Real* e, + const T* X, + const T* Y, + const int& dim, + const int& nst, + const int& prec = 8) -> void { + if (GlobalV::MY_RANK == 0) + { + assert(nst == LR_Util::write_value(efile_out("full_"+label), prec, e, nst)); + } + assert(nst * dim == LR_Util::write_value(vfile_out("full_X_"+label), prec, X, nst, dim)); + assert(nst * dim == LR_Util::write_value(vfile_out("full_Y_"+label), prec, Y, nst, dim)); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "write full states " + label); + }; + + if ((this->input.bse_tda == "both" || this->input.bse_tda == "tda")) + { + for (int is = 0; is < this->input.bse_spin_types.size(); ++is) + { + bse_matrix.tda_solver(is, + this->nstates, + &this->tda_ene[is * this->nstates], + this->X[is].template data()); + + std::cout << "eigenvalues: (Ry)" << std::endl; + int write_nstates = std::min(this->nstates, 20); + LR_Util::print_value(&this->tda_ene[is * this->nstates], write_nstates); + std::cout << "eigenvalues: (eV)" << std::endl; + for (int i = 0; i < write_nstates; ++i) + { + std::cout << this->tda_ene[is * this->nstates + i] * ModuleBase::Ry_to_eV << " "; + } + std::cout << std::endl; + std::cout << "Excition binding energies (eV):" + << (direct_gap - tda_ene[is * this->nstates]) * ModuleBase::Ry_to_eV << std::endl; + + if (this->input.out_wfc_lr) + { + write_tda_states(this->input.bse_spin_types[is], + &this->tda_ene[is * this->nstates], + this->X[is].template data(), + this->nloc_per_state, + this->nstates); + } + malloc_trim(0); + } + } + if ((this->input.bse_tda == "both" || this->input.bse_tda == "full")) + { + for (int is = 0; is < this->input.bse_spin_types.size(); ++is) + { + bse_matrix.full_solver(is, + this->nstates, + &this->full_ene[is * this->nstates], + this->full_X[is].template data(), + this->full_Y[is].template data()); + + std::cout << "eigenvalues: (Ry)" << std::endl; + int write_nstates = std::min(this->nstates, 20); + LR_Util::print_value(&this->full_ene[is * this->nstates], write_nstates); + for (int i = 0; i < write_nstates; ++i) + { + std::cout << this->full_ene[is * this->nstates + i] * ModuleBase::Ry_to_eV << " "; + } + std::cout << std::endl; + std::cout << "Excition binding energies (eV):" + << (direct_gap - full_ene[is * this->nstates]) * ModuleBase::Ry_to_eV << std::endl; + + if (this->input.out_wfc_lr) + { + write_full_states(this->input.bse_spin_types[is], + &this->full_ene[is * this->nstates], + this->full_X[is].template data(), + this->full_Y[is].template data(), + this->nloc_per_state, + this->nstates); + } + malloc_trim(0); + } + } + } + } + else if (this->input.lr_solver == "spectrum" || this->input.lr_solver == "plot") + { + std::cout << "Reading BSE excitation states from file." << std::endl; + auto read_tda_states = [&](const std::string& label, Real* e, T* v, const int& dim, const int& nst)->void + { + if (GlobalV::MY_RANK == 0) { + assert(nst == LR_Util::read_value(efile_in(label), e, nst)); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "finish reading " + efile_in(label)); + } +#ifdef __MPI + MPI_Bcast(e, nst, MPI_DOUBLE, 0, MPI_COMM_WORLD); +#endif + assert(nst * dim == LR_Util::read_value(vfile_in(label), v, nst, dim)); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "finish reading " + vfile_in(label)); + }; + auto read_full_states = [&](const std::string& label, Real* e, T* X, T* Y, const int& dim, const int& nst)->void + { + if (GlobalV::MY_RANK == 0) { + assert(nst == LR_Util::read_value(efile_in("full_"+label), e, nst)); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "finish reading " + efile_in("full_"+label)); + } +#ifdef __MPI + MPI_Bcast(e, nst, MPI_DOUBLE, 0, MPI_COMM_WORLD); +#endif + assert(nst * dim == LR_Util::read_value(vfile_in("full_X_"+label), X, nst, dim)); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "finish reading " + vfile_in("full_X_"+label)); + assert(nst * dim == LR_Util::read_value(vfile_in("full_Y_"+label), Y, nst, dim)); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "finish reading " + vfile_in("full_Y_"+label)); + }; + + if (this->input.bse_tda == "both" || this->input.bse_tda == "tda") + { + for (int is = 0; is < this->input.bse_spin_types.size(); ++is) + { + read_tda_states(this->input.bse_spin_types[is], + &this->tda_ene[is * this->nstates], + this->X[is].template data(), + this->nloc_per_state, + this->nstates); + } + } + if (this->input.bse_tda == "both" || this->input.bse_tda == "full") + { + for (int is = 0; is < this->input.bse_spin_types.size(); ++is) + { + read_full_states(this->input.bse_spin_types[is], + &this->full_ene[is * this->nstates], + this->full_X[is].template data(), + this->full_Y[is].template data(), + this->nloc_per_state, + this->nstates); + // check whether |X|^2 - |Y|^2 = 1 + for (int i = 0; i < this->nstates; ++i) + { + double norm_xy = 0.0; + for (int j = 0; j < this->nloc_per_state; ++j) { + norm_xy += std::norm(this->full_X[is].template data()[i * this->nloc_per_state + j]) + - std::norm(this->full_Y[is].template data()[i * this->nloc_per_state + j]); + } + Parallel_Reduce::reduce_all(norm_xy); + if (std::abs(norm_xy - 1.0) > 1e-6){ + std::cout << "| CHECK WARNING: for full excitation " << i + << ", |X|^2 - |Y|^2 = " << std::setprecision(10) << norm_xy << std::endl; + } + } + } + } + } + else + { + ModuleBase::WARNING_QUIT("ESolver_BSE", "lr_solver must be elpa, plot or spectrum"); + } + ModuleBase::timer::end("ESolver_BSE", "runner"); + return; +} + +template +void ESolver_BSE::after_all_runners(UnitCell& ucell) +{ + ModuleBase::TITLE("ESolver_BSE", "after_all_runners"); + ModuleBase::timer::start("ESolver_BSE", "after_all_runners"); + const std::string& output_dir = PARAM.globalv.global_out_dir; + const std::set benchmarks = {"abacus-librpa", "abacus", "none" }; + if (benchmarks.find(this->input.ri_hartree_benchmark) == benchmarks.end()) + { + return; + } // no need to calculate the spectrum + + if (this->input.lr_solver == "plot") + { + for (int is = 0; is < this->X.size(); ++is) + { + std::cout << "plot BSE exciton wavefunction for state: " << this->input.plot_istate + << ", spin type: " << this->input.bse_spin_types[is] << std::endl; + LR_Util::ExcitonPlotter eplot(this->nspin, this->nbasis, this->nocc, this->nvirt, *this->psi_ks, + this->ucell, this->kv, this->gd, this->orb_cutoff_, this->Pgrid, *this->pw_rho, + this->paraX_, this->paraC_, this->paraMat_, + output_dir, + &this->tda_ene[is * this->nstates], this->X[is].template data(), + false/*openshell*/, &this->orb_); + const std::string plot_type = LR_Util::tolower(this->input.exciton_plot_type); + const std::string plot_format = LR_Util::tolower(this->input.exciton_plot_format); + const bool write_slice = (plot_format == "slice" || plot_format == "both"); + const bool write_cube = (plot_format == "cube" || plot_format == "both"); + + if (plot_format != "cube" && plot_format != "slice" && plot_format != "both") + { + ModuleBase::WARNING_QUIT("ESolver_BSE", "exciton_plot_format must be cube, slice, or both"); + } + + if (plot_type == "conditional") + { + if (this->input.exciton_fixed_coordinate.size() != 6) + { + ModuleBase::WARNING_QUIT( + "ESolver_BSE", + "exciton_fixed_coordinate must contain six values: hole x y z followed by electron x y z"); + } + const std::array r_h_fix = {this->input.exciton_fixed_coordinate[0], + this->input.exciton_fixed_coordinate[1], + this->input.exciton_fixed_coordinate[2]}; + const std::array r_e_fix = {this->input.exciton_fixed_coordinate[3], + this->input.exciton_fixed_coordinate[4], + this->input.exciton_fixed_coordinate[5]}; + if (write_cube) + { + eplot.plot_conditional_density(this->input.plot_istate, r_h_fix, "elec"); + eplot.plot_conditional_density(this->input.plot_istate, r_e_fix, "hole"); + } + if (write_slice) + { + eplot.plot_cond_slice(this->input.plot_istate, r_h_fix, + this->input.exciton_slice_plane, + this->input.exciton_slice_pos, + this->input.exciton_slice_npoints, + this->input.exciton_slice_scale, "elec"); + eplot.plot_cond_slice(this->input.plot_istate, r_e_fix, + this->input.exciton_slice_plane, + this->input.exciton_slice_pos, + this->input.exciton_slice_npoints, + this->input.exciton_slice_scale, "hole"); + } + } + else if (plot_type == "average") + { + if (write_cube) + { + // Average hole density: integrates out the electron coordinate + eplot.plot_average_density(this->input.plot_istate, "hole"); + // Average electron density: integrates out the hole coordinate + eplot.plot_average_density(this->input.plot_istate, "elec"); + } + if (write_slice) + { + eplot.plot_average_slice(this->input.plot_istate, "hole", + this->input.exciton_slice_plane, + this->input.exciton_slice_pos, + this->input.exciton_slice_npoints, + this->input.exciton_slice_scale); + eplot.plot_average_slice(this->input.plot_istate, "elec", + this->input.exciton_slice_plane, + this->input.exciton_slice_pos, + this->input.exciton_slice_npoints, + this->input.exciton_slice_scale); + } + } + else + { + ModuleBase::WARNING_QUIT("ESolver_BSE", "exciton_plot_type must be average or conditional"); + } + } + } + if (this->input.lr_solver == "spectrum" || this->input.lr_solver == "elpa") + { + std::cout << "Calculating BSE optical absorption spectrum." << std::endl; + if (LR_Util::tolower(this->input.abs_gauge) == "velocity" ) + { + const int nspin_tmp = PARAM.inp.nspin == 2 ? 2 : 1; + this->velocity_mo = LR_Util::cal_velocity_mo(this->ucell, this->gd, this->two_center_bundle_, + this->paraMat_, this->paraC_, this->kv, *this->psi_ks, + this->nk, nspin_tmp, this->nbasis, this->nocc, this->nvirt); + } + if (this->input.bse_tda == "both" || this->input.bse_tda == "tda") + { + for (int is = 0; is < this->X.size(); ++is) + { + LR::LR_Spectrum spectrum(this->nspin, this->nbasis, this->nocc, this->nvirt, *this->pw_rho, *this->psi_ks, + this->ucell, this->kv, this->gd, this->orb_cutoff_, this->two_center_bundle_, + this->paraX_, this->paraC_, this->paraMat_, + &this->tda_ene[is * this->nstates], this->eig_ks.c, + this->X[is].template data(), this->nstates, false/*openshell*/, + LR_Util::tolower(this->input.abs_gauge)); + if (LR_Util::tolower(this->input.abs_gauge) == "velocity") + { + spectrum.set_vmo(this->velocity_mo.data()); + } + spectrum.cal_spectrum(); + spectrum.transition_analysis(this->input.bse_spin_types[is]+"_tda"); + if (this->input.bse_spin_types[is] != "triplet") // triplets has no transition dipole and no contribution to the spectrum + { + spectrum.write_transition_dipole(output_dir + + "trans_dipole_" + this->input.bse_spin_types[is] + "_tda.dat"); + // ============================== for test ============================== + if (LR_Util::tolower(this->input.abs_gauge) == "velocity") + { //// TEST the formula v/omega rather than v/(e_a-e_i) + // spectrum.test_transition_dipoles_velocity_omega(); + // spectrum.write_transition_dipole(PARAM.globalv.global_out_dir + + // "trans_dipole_" + spin_types[is] + "_vomega_tda.dat"); + } + // ============================== for test ============================== + } + } + } + if (this->input.bse_tda == "both" || this->input.bse_tda == "full") + { + for (int is = 0;is < this->full_X.size();++is) + { + LR::LR_Spectrum spectrum(this->nspin, this->nbasis, this->nocc, this->nvirt, *this->pw_rho, *this->psi_ks, + this->ucell, this->kv, this->gd, this->orb_cutoff_, this->two_center_bundle_, + this->paraX_, this->paraC_, this->paraMat_, + &this->full_ene[is * this->nstates], this->eig_ks.c, + this->full_X[is].template data(), this->nstates, false/*openshell*/, + LR_Util::tolower(this->input.abs_gauge)); + if (LR_Util::tolower(this->input.abs_gauge) == "velocity") + { + spectrum.set_vmo(this->velocity_mo.data()); + } + spectrum.set_Y(this->full_Y[is].template data()); + spectrum.set_full(true); + spectrum.cal_spectrum(); + spectrum.transition_analysis(this->input.bse_spin_types[is]+"_full"); + if (this->input.bse_spin_types[is] != "triplet") // triplets has no transition dipole and no contribution to the spectrum + { + spectrum.write_transition_dipole(output_dir + + "trans_dipole_" + this->input.bse_spin_types[is] + "_full.dat"); + } + } + } + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "ESolver_BSE::after_all_runners"); + ModuleBase::timer::end("ESolver_BSE", "after_all_runners"); +} + +template +void ESolver_BSE::ipa_solver() +{// if ipa, assign X as identity matrix directly + ModuleBase::TITLE("ESolver_BSE", "ipa_solver"); + ModuleBase::timer::start("ESolver_BSE", "ipa_solver"); + std::cout << "Independent particle approximation is used, assign X as identity matrix directly." << std::endl; + assert(this->input.bse_tda == "tda"); + std::vector ev(this->nk * this->nocc[0] * this->nvirt[0], 0.0); + for (int ik = 0; ik < this->nk; ++ik) + { + for (int i = 0; i < this->nocc[0]; ++i) + { + for (int a = 0; a < this->nvirt[0]; ++a) + { + int index = ik * this->nocc[0] * this->nvirt[0] + i * this->nvirt[0] + a; + ev[index] = this->eig_gw(ik, this->nocc[0] + a) - this->eig_gw(ik, i); + } + } + } + + std::vector indices(ev.size()); + std::iota(indices.begin(), indices.end(), 0); // [0, 1, 2, ..., size-1] + + std::sort(indices.begin(), indices.end(), [&](int lhs, int rhs) { + return ev[lhs] < ev[rhs]; + }); + std::sort(ev.begin(), ev.end()); + std::copy_n(ev.data(), this->nstates, this->tda_ene.data()); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::size_t istate = 0; istate < this->nstates; ++istate) + { + int sorted_index = indices[istate]; + int ik = sorted_index / (this->nocc[0] * this->nvirt[0]); + int loffset_X = (istate * this->nk + ik) * this->paraX_[0].get_local_size(); + int i = (sorted_index / this->nvirt[0]) % this->nocc[0]; + int a = sorted_index % this->nvirt[0]; + int col_loc = this->paraX_[0].global2local_col(i); + int row_loc = this->paraX_[0].global2local_row(a); + if (col_loc == -1 || row_loc == -1) continue; + this->X[0].template data()[loffset_X + col_loc * this->paraX_[0].get_row_size() + row_loc] = 1.0; + } + ModuleBase::timer::end("ESolver_BSE", "ipa_solver"); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "IPA solver"); +} + +template +void ESolver_BSE::lri_init() +{ + ModuleBase::TITLE("ESolver_BSE", "LRI init"); + using TA = int; + using TC = std::array; + using TAC = std::pair; + // start reading Ws and Cs + std::map>> Cs_in; + std::map>> Vs_in; + std::map>> Ws_in; + const std::string& dir = PARAM.globalv.global_readin_dir; + // if (GlobalV::MY_RANK == 0) // comment to read from all processes to avoid communication + // { + Cs_in = LRI_CV_Tools::read_Cs_ao_all(dir); + if (this->input.ri_hartree_benchmark == "aims-librpa" ) + { + Vs_in = LR_IO::read_coulomb_mat_general_k(dir, Cs_in, this->kRlist); + } + else if (this->input.ri_hartree_benchmark == "none" || this->input.ri_hartree_benchmark == "abacus-librpa" ) + { + Vs_in = LR_IO::read_coulomb_mat_k(dir, Cs_in, this->kRlist); + } + Ws_in = LR_IO::read_Ws(Vs_in, this->kRlist.Rlist); + // } +#ifdef __MPI + MPI_Barrier(MPI_COMM_WORLD); +#endif + this->mo_lri->init(Cs_in, Vs_in, Ws_in, GlobalC::exx_info.info_ri); + malloc_trim(0); + ModuleBase::TITLE("ESolver_BSE", "Finish LRI init"); +} + +template +void ESolver_BSE::read_ks_wfc() +{ + assert(this->psi_ks != nullptr); + this->pelec->ekb.create(this->kv.get_nks(), this->nbands); + this->pelec->wg.create(this->kv.get_nks(), this->nbands); + this->eig_gw.create(this->kv.get_nks(), this->nbands); + + int ncore = 0; // skip core bands + int nbands_file = 0; + int nk_file = 0; + int nspin_file = 0; + int nocc_file = 0; + int nspin_tmp = PARAM.inp.nspin == 2 ? 2 : 1; + LR_IO::parse_band_out_file(nbands_file, nk_file, nspin_file, nocc_file); + if (nk_file != this->nk) { + ModuleBase::WARNING_QUIT("ESolver_BSE", "Inconsistence: The nk in `band_out` is " + std::to_string(nk_file) + + ", while BSE::nk is " + std::to_string(this->nk)); + } + std::vector eig_gw_info; + if (PARAM.inp.bse_use_fine_kgrid) + { + eig_gw_info = LR_IO::read_energy_qp_from_band_files(this->kv, this->nocc[0], this->nvirt[0], ncore, + this->nk, nspin_tmp, nspin_file); + LR_IO::read_librpa_eigenvectors_from_band_files(*this->psi_ks, *this->psi_ks_global, + PARAM.globalv.global_readin_dir, ncore, nbands_file, nspin_tmp, nspin_file, this->paraMat_); + } + else + { + eig_gw_info = LR_IO::read_energy_qp(this->nocc[0], this->nvirt[0], + ncore, this->nk, nspin_tmp, nspin_file); + LR_IO::read_librpa_eigenvectors(*this->psi_ks, *this->psi_ks_global, + PARAM.globalv.global_readin_dir, ncore, nbands_file, nspin_tmp, nspin_file, this->paraMat_); + } + int cbm_k(0), vbm_k(0), direct_k(0); + for (int iks = 0; iks < this->kv.get_nks(); ++iks) { + for (int ib = 0; ib < this->nbands; ++ib) { + this->pelec->wg(iks, ib) = eig_gw_info[iks * this->nbands *3 + ib * 3 + 0]; + this->pelec->ekb(iks, ib) = eig_gw_info[iks * this->nbands *3 + ib * 3 + 1]; + this->eig_gw(iks, ib) = eig_gw_info[iks * this->nbands *3 + ib * 3 + 2]; + } + double cbm = this->eig_gw(iks, this->nocc[0]); + for (int ib = this->nocc[0]; ib < this->nbands; ++ib) { // in case of non-ordered bands + double e = this->eig_gw(iks, ib); + if (e < cbm) cbm = e; + } + double vbm = this->eig_gw(iks, this->nocc[0]-1); + for (int ib = 0; ib < this->nocc[0]-1; ++ib) { + double e = this->eig_gw(iks, ib); + if (e > vbm) vbm = e; + } + if (iks == 0) { + this->cbm_energy = cbm; + this->vbm_energy = vbm; + this->direct_gap = cbm - vbm; + } + else { + if (this->cbm_energy > cbm) { + this->cbm_energy = cbm; + cbm_k = iks; + } + if (this->vbm_energy < vbm) { + this->vbm_energy = vbm; + vbm_k = iks; + } + if (this->direct_gap > cbm - vbm) { + this->direct_gap = cbm - vbm; + direct_k = iks; + } + } + } + std::cout << "VBM energy (eV): " << this->vbm_energy * ModuleBase::Ry_to_eV << " at k " << vbm_k << std::endl; + std::cout << "CBM energy (eV): " << this->cbm_energy * ModuleBase::Ry_to_eV << " at k " << cbm_k << std::endl; + std::cout << "Indirect gap (eV): " << (this->cbm_energy - this->vbm_energy) * ModuleBase::Ry_to_eV << std::endl; + std::cout << "Direct gap (eV): " << this->direct_gap * ModuleBase::Ry_to_eV << " at k " << direct_k << std::endl; + + this->eig_ks = std::move(this->pelec->ekb); +} + +template +void ESolver_BSE::init_pot(const Charge& chg_gs) +{ + switch (this->nspin) + { + using ST = LR::PotHxcLR::SpinType; + case 1: case 2: + this->pot[0] = std::make_shared(this->xc_kernel, *this->pw_rho, this->ucell, chg_gs, this->Pgrid, + ST::S1, this->input.lr_init_xc_kernel); + break; + // case 2: + // this->pot[0] = std::make_shared(xc_kernel, *this->pw_rho, ucell, chg_gs, Pgrid, openshell ? ST::S2_updown : ST::S2_singlet, input.lr_init_xc_kernel); + // this->pot[1] = std::make_shared(xc_kernel, *this->pw_rho, ucell, chg_gs, Pgrid, openshell ? ST::S2_updown : ST::S2_triplet, input.lr_init_xc_kernel); + // break; + default: + throw std::invalid_argument("ESolver_BSE: nspin must be 1 or 2"); + } +} + +template +void ESolver_BSE::allocate_eigen_infos() +{ + ModuleBase::TITLE("ESolver_BSE", "allocate_eigen_infos"); + + for (int is = 0; is < this->nspin; ++is) + { + Parallel_2D px; + LR_Util::setup_2d_division(px, /*nb2d=*/1, this->nvirt[is], this->nocc[is] +#ifdef __MPI + , this->paraC_.blacs_ctxt +#endif + ); + this->paraX_.emplace_back(std::move(px)); + } + this->nloc_per_state = this->nk + * (this->openshell ? this->paraX_[0].get_local_size() + this->paraX_[1].get_local_size() + : this->paraX_[0].get_local_size()); + + int n_spin_types = this->input.bse_spin_types.size(); + if (this->input.bse_tda == "both" || this->input.bse_tda == "tda") { + BSE_Util::print_mem_estimate("TDA BSE eigen states", + n_spin_types * static_cast(this->nstates) + * (1 + this->nloc_per_state), + sizeof(T)); + this->tda_ene.resize(n_spin_types * this->nstates); + this->X.resize(n_spin_types, LR_Util::newTensor({ this->nstates, this->nloc_per_state })); + for (auto& x : this->X) { x.zero(); } + } + if (this->input.bse_tda == "both" || this->input.bse_tda == "full") { + BSE_Util::print_mem_estimate("full BSE eigen states", + n_spin_types * static_cast(this->nstates) + * (1 + 2 * this->nloc_per_state), + sizeof(T)); + this->full_ene.resize(n_spin_types * this->nstates); + this->full_X.resize(n_spin_types, LR_Util::newTensor({ this->nstates, this->nloc_per_state })); + this->full_Y.resize(n_spin_types, LR_Util::newTensor({ this->nstates, this->nloc_per_state })); + for (auto& x : this->full_X) { x.zero(); } + for (auto& y : this->full_Y) { y.zero(); } + } +} + +template class ESolver_BSE; +template class ESolver_BSE, double>; +} // namespace BSE diff --git a/source/source_lcao/module_lr/bse/esolver_bse_lcao.h b/source/source_lcao/module_lr/bse/esolver_bse_lcao.h new file mode 100644 index 0000000000..52767f7555 --- /dev/null +++ b/source/source_lcao/module_lr/bse/esolver_bse_lcao.h @@ -0,0 +1,57 @@ +#pragma once +#include "hamilt_bse.h" +#include "source_io/module_parameter/parameter.h" +#include "source_lcao/module_lr/esolver_lrtd_lcao.h" +#include "source_lcao/module_lr/lr_spectrum.h" +#include "source_lcao/module_lr/ri_benchmark/ri_benchmark.h" +#include "source_lcao/module_lr/utils/lr_io.h" +#include "source_lcao/module_lr/utils/exciton_plotter.h" +#include "source_lcao/module_ri/LRI_CV_Tools.h" + +namespace BSE +{ + template using Real = typename GetTypeReal::type; + + template + class ESolver_BSE : public LR::ESolver_LR { + public: + /// @brief a from-scratch constructor + ESolver_BSE(const Input_para& inp, UnitCell& ucell); + + ~ESolver_BSE() override { + //delete this->psi_ks; // already deleted in ESolver_LR + delete this->psi_ks_global; + } + + LR_IO::RI_kRlist kRlist; + psi::Psi* psi_ks_global; ///< global version of psi_ks + ModuleBase::matrix eig_gw; ///< GW energy + double cbm_energy, vbm_energy; //conduction band minimum and valence band maximum + double direct_gap; + std::vector tda_ene, full_ene; // in Rydberg + + /// @brief [nspin_types][{nstates, nk* (locc* lvirt}] + std::vector full_X, full_Y; + + std::unique_ptr> mo_lri; + + virtual void runner(UnitCell& ucell, int istep) override; + virtual void after_all_runners(UnitCell& ucell) override; + + void lri_init(); + + /// @brief solve IPA without construct and diagonalize matrix + void ipa_solver(); + + /// @brief read in the ground state wave function, gw band energy and occupation + virtual void read_ks_wfc() override; + + /// @brief init Hartree potential for BSE, + /// @attention here we don't multiply 2 for singlet or 0 for triplet, we will do it in HamiltBSE + virtual void init_pot(const Charge& chg_gs) override; + + /// @brief X for tda and also Y for full BSE excitation + void allocate_eigen_infos(); + }; + +} \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/hamilt_bse.cpp b/source/source_lcao/module_lr/bse/hamilt_bse.cpp new file mode 100644 index 0000000000..f0a62df7c5 --- /dev/null +++ b/source/source_lcao/module_lr/bse/hamilt_bse.cpp @@ -0,0 +1,603 @@ +#include "hamilt_bse.h" + +#include "hamilt_bse_solver.h" +#include "source_lcao/module_gint/gint_interface.h" +#include "source_lcao/module_lr/utils/lr_util_hcontainer.h" + +#include + +namespace BSE +{ +template +HamiltBSE::HamiltBSE(const int& nspin, + const int& naos, + const std::vector& nocc, + const std::vector& nvirt, + const UnitCell& ucell_in, + const std::vector& orb_cutoff_in, + const Grid_Driver& gd_in, + const psi::Psi& psi_in, + const psi::Psi& psi_glb_in, + const ModuleBase::matrix& eig_gw_in, + MolecularLRI& mo_lri_in, + std::weak_ptr pot_in, + const K_Vectors& kv_in, + const std::vector& pX_in,// vector for spin, parallel as {nvirt, nocc} + const Parallel_2D& pc_in, //parallel as {nbasis, nbands} + const Parallel_Orbitals& pmat_in, //parallel as {nbasis, nbasis} + const std::vector& spin_types_in, //can be singlet and triplet + const std::string& tda, // can be: "tda", "full", "both" + const std::string& ri_hartree_benchmark_in) + : nspin(nspin), naos(naos), nocc(nocc), nvirt(nvirt), ucell(ucell_in), + orb_cutoff(orb_cutoff_in), gd(gd_in), psi_ks(psi_in), psi_ks_glb(psi_glb_in), eig_gw(eig_gw_in), + mo_lri(mo_lri_in), + pot(pot_in), kv(kv_in), + pX(pX_in), pc(pc_in), pmat(pmat_in), + spin_types(spin_types_in), ri_hartree_benchmark(ri_hartree_benchmark_in) +{ + ModuleBase::TITLE("BSE", "HamiltBSE"); + if (this->pX[0].get_local_size() == 0) { + std::cerr<< "Warning: Parallel_2D in RANK "+std::to_string(GlobalV::MY_RANK) +" has no local size, please use less mpi." << std::endl; + std::cerr<< " [File:"<<__FILE__<< ", Function: " << __FUNCTION__ << ", Line: " << __LINE__ << "]" << std::endl; + } + assert(naos == pmat.get_global_row_size() && naos == pmat.get_global_col_size()); + this->nk = this->nspin == 2 ? this->kv.get_nks() / 2 : this->kv.get_nks(); + this->ndim = nk * nocc[0] * nvirt[0]; + + int nb2d; + if (this->ndim > 1000) + { + nb2d = 64; + } + else if (this->ndim > 500) + { + nb2d = 32; + } + else if (this->ndim > 0) + { + nb2d = 1; + } + else throw std::runtime_error("ndim in HamiltBSE is zero or negative."); + LR_Util::setup_2d_division(this->pA, nb2d, ndim, ndim + #ifdef __MPI + , this->pX[0].blacs_ctxt + #endif + ); + + if (!PARAM.inp.bse_ri_hartree && this->ri_hartree_benchmark == "none") + { + this->DM_trans = LR_Util::make_unique>(&pmat, 1/*nspin*/, kv_in.kvec_d, nk); + this->DM_trans->set_DMK_zero(); + LR_Util::initialize_DMR(*this->DM_trans, this->pmat, this->ucell, this->gd, this->orb_cutoff); + } + if (PARAM.inp.bse_mem_save) { assert(PARAM.inp.bse_continue == 0 && PARAM.inp.bse_ri_hartree); } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "HamiltBSE is ready to calculate"); + + const std::string read_dir = PARAM.globalv.global_readin_dir; + if (PARAM.inp.bse_continue >= 1) { + BSE_Util::print_mem_estimate("V matrix of A", this->pA.get_local_size(), sizeof(T)); + this->VA_local.resize(this->pA.get_local_size(), 0.0); + this->read_AB_matrix(read_dir + "A_V_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", this->VA_local.data(), this->ndim, this->ndim); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read_V_for_A"); + } + if (PARAM.inp.bse_continue >= 2) { + BSE_Util::print_mem_estimate("W matrix of A", this->pA.get_local_size(), sizeof(T)); + this->WA_local.resize(this->pA.get_local_size(), 0.0); + this->read_AB_matrix(read_dir + "A_W_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", this->WA_local.data(), this->ndim, this->ndim); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read_W_for_A"); + } + if (PARAM.inp.bse_continue >= 3) { + BSE_Util::print_mem_estimate("V matrix of B", this->pA.get_local_size(), sizeof(T)); + this->VB_local.resize(this->pA.get_local_size(), 0.0); + this->read_AB_matrix(read_dir + "B_V_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", this->VB_local.data(), this->ndim, this->ndim); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read_V_for_B"); + } + if (PARAM.inp.bse_continue >= 4) { + BSE_Util::print_mem_estimate("W matrix of B", this->pA.get_local_size(), sizeof(T)); + this->WB_local.resize(this->pA.get_local_size(), 0.0); + this->read_AB_matrix(read_dir + "B_W_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", this->WB_local.data(), this->ndim, this->ndim); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read_W_for_B"); + } + + if (!PARAM.inp.bse_mem_save) + { + for (const auto& st : this->spin_types) { + if (st == "singlet" || st == "triplet") { + // Hartree term V (exchange electron and hole) + if (st == "singlet"){ + this->cal_V_for_A(); + if (tda == "both" || tda == "full") { this->cal_V_for_B(); } + } + else if (st == "triplet") { + std::cout << "Hatree term is not needed for triplet." << std::endl; + } + + // direct term W (electron-electron and hole-hole) + this->cal_W_for_A(); + if (tda == "both" || tda == "full") { this->cal_W_for_B(); } + } + else if(st == "rpa") { + this->cal_V_for_A(); + if (tda == "both" || tda == "full") { this->cal_V_for_B(); } + } + else if(st != "ipa") { + throw std::runtime_error("Unsupported type in BSE: " + st); + } + } + this->mo_lri.LR_lri.free_Vs(); + this->mo_lri.LR_lri.free_Ws(); + malloc_trim(0); + } +} + +template +void HamiltBSE::cal_V_for_A(){ + ModuleBase::TITLE("HamiltBSE", "cal_V_for_A"); + ModuleBase::timer::start("HamiltBSE", "cal_V_for_A"); + std::cout<<"in cal_V_for_A"<VA_local.empty()) { + std::cout<< "V for A has been calculated, skip." <pA.get_local_size(), sizeof(T)); + this->VA_local.resize(this->pA.get_local_size(), 0.0); + if (this->ri_hartree_benchmark == "aims" || this->ri_hartree_benchmark == "abacus") { + throw std::runtime_error("this BSE routine only supports aims-librpa/abacus-librpa benchmark"); + } + else if (PARAM.inp.bse_ri_hartree || this->ri_hartree_benchmark =="aims-librpa" || this->ri_hartree_benchmark == "abacus-librpa") { + std::cout << "Calculating Hartree term for A with RI approximation" << std::endl; + this->mo_lri.cal_hartree_for_A(this->VA_local, this->pA); + } + else if (this->ri_hartree_benchmark == "none") { // do things like OperatorLRHxc + std::cout << "Calculating Hartree term for A with grid integration" << std::endl; + this->cal_V_by_grid(true); + } + if (PARAM.inp.out_bse_ab){ + this->write_AB_matrix(PARAM.globalv.global_out_dir+"A_V_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", 6, this->VA_local.data(), this->ndim, this->ndim); + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "cal_V_for_A"); + ModuleBase::timer::end("HamiltBSE", "cal_V_for_A"); +} + +template +void HamiltBSE::cal_V_for_B(){ + ModuleBase::TITLE("HamiltBSE", "cal_V_for_B"); + ModuleBase::timer::start("HamiltBSE", "cal_V_for_B"); + std::cout<<"in cal_V_for_B"<VB_local.empty()) { + std::cout<< "V for B has been calculated, skip." <pA.get_local_size(), sizeof(T)); + this->VB_local.resize(this->pA.get_local_size(), 0.0); + if (this->ri_hartree_benchmark == "aims" || this->ri_hartree_benchmark == "abacus") { + throw std::runtime_error("this BSE routine only supports aims-librpa/abacus-librpa benchmark"); + } + else if (PARAM.inp.bse_ri_hartree || this->ri_hartree_benchmark =="aims-librpa" || this->ri_hartree_benchmark == "abacus-librpa") { + std::cout << "Calculating Hartree term for B with RI approximation" << std::endl; + this->mo_lri.cal_hartree_for_B(this->VB_local, this->pA); + } + else if (this->ri_hartree_benchmark == "none") { // do things like OperatorLRHxc + std::cout << "Calculating Hartree term for B with grid integration" << std::endl; + this->cal_V_by_grid(false); + } + if (PARAM.inp.out_bse_ab){ + this->write_AB_matrix(PARAM.globalv.global_out_dir+"B_V_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", 6, this->VB_local.data(), this->ndim, this->ndim); + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "cal_V_for_B"); + ModuleBase::timer::end("HamiltBSE", "cal_V_for_B"); +} + +template +void HamiltBSE::cal_W_for_A(){ + ModuleBase::TITLE("HamiltBSE", "cal_W_for_A"); + ModuleBase::timer::start("HamiltBSE", "cal_W_for_A"); + std::cout<<"in cal_W_for_A"<WA_local.empty()) { + std::cout<< "W for A has been calculated, skip." <pA.get_local_size(), sizeof(T)); + this->WA_local.resize(this->pA.get_local_size(), 0.0); + this->mo_lri.cal_W_for_A(this->WA_local, this->pA); + if (PARAM.inp.out_bse_ab){ + this->write_AB_matrix(PARAM.globalv.global_out_dir+"A_W_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", 6, this->WA_local.data(), this->ndim, this->ndim); + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "cal_W_for_A"); + ModuleBase::timer::end("HamiltBSE", "cal_W_for_A"); +} + +template +void HamiltBSE::cal_W_for_B(){ + ModuleBase::TITLE("HamiltBSE", "cal_W_for_B"); + ModuleBase::timer::start("HamiltBSE", "cal_W_for_B"); + std::cout<<"in cal_W_for_B"<WB_local.empty()) { + std::cout<< "W for B has been calculated, skip." <pA.get_local_size(), sizeof(T)); + this->WB_local.resize(this->pA.get_local_size(), 0.0); + this->mo_lri.cal_W_for_B(this->WB_local, this->pA); + + if (PARAM.inp.out_bse_ab){ + this->write_AB_matrix(PARAM.globalv.global_out_dir+"B_W_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", 6, this->WB_local.data(), this->ndim, this->ndim); + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "cal_W_for_B"); + ModuleBase::timer::end("HamiltBSE", "cal_W_for_B"); +} + + +template +void HamiltBSE::init_bse_matrix(const bool is_full, const int & st_index){ + ModuleBase::TITLE("HamiltBSE", "init_bse_matrix"); + ModuleBase::timer::start("HamiltBSE", "init_bse_matrix"); + + double alpha, beta; + const std::string& st = this->spin_types[st_index]; + if (st == "singlet") { alpha = 2.0; beta = -1.0; } + else if (st == "triplet") { alpha = 0.0; beta = -1.0; } + else if (st == "rpa") { alpha = 2.0; beta = 0.0; } + else if (st == "ipa") { alpha = 0.0; beta = 0.0; } + else { throw std::runtime_error("Unsupported type in BSE: " + st); } + + std::string tda_type = is_full ? "full" : "TDA"; + std::cout<<"| init "<< tda_type << " BSE for type: "<< st << std::endl; + std::cout<<"| A(ai,bj) = (Ea-Ei) δ_ij δ_ab + α (ai|V|jb) + β (ji|W|ab)" << std::endl; + std::cout<<"| term coefficient: (Exchange) α: "<pA.get_local_size(), sizeof(T)); + this->BSE_A_local.assign(this->pA.get_local_size(), 0.0); + if (is_full) + { + BSE_Util::print_mem_estimate("BSE B matrix", this->pA.get_local_size(), sizeof(T)); + this->BSE_B_local.assign(this->pA.get_local_size(), 0.0); + } + // 1) add diagonal GW energy term +#ifdef _OPENMP +#pragma omp parallel for collapse(3) +#endif + for(int ik = 0;ik < nk;++ik) + { + for (int i = 0;i < nocc[0];++i) + { + for(int a = 0;a < nvirt[0];++a) + { + int index = ik * nocc[0] * nvirt[0] + i * nvirt[0] + a; + int col_loc = this->pA.global2local_col(index); + int row_loc = this->pA.global2local_row(index); + if (col_loc == -1 || row_loc == -1) continue; + this->BSE_A_local[col_loc * this->pA.get_row_size() + row_loc] + = this->eig_gw(ik, nocc[0] + a) - this->eig_gw(ik, i); + } + } + } + // 2) add V/W contributions + if (PARAM.inp.bse_mem_save) + { + std::cout << "| bse_mem_save is true, V and W matrix will be added to BSE matrix directly." << std::endl; + if (alpha != 0.0) { + this->mo_lri.cal_hartree_for_A(this->BSE_A_local, this->pA, alpha); + if (is_full) { this->mo_lri.cal_hartree_for_B(this->BSE_B_local, this->pA, alpha); } + } + if (beta != 0.0) { + this->mo_lri.cal_W_for_A(this->BSE_A_local, this->pA, beta); + if (is_full) { this->mo_lri.cal_W_for_B(this->BSE_B_local, this->pA, beta); } + } + GlobalV::ofs_running << "| V and W matrix has been added." << std::endl; + std::cout << "| V and W matrix has been added." << std::endl; + } + else + { + if (alpha != 0.0) { + assert(this->VA_local.size() == this->BSE_A_local.size()); + if (is_full) assert(this->VB_local.size() == this->BSE_A_local.size()); + } + if (beta != 0.0) { + assert(this->WA_local.size() == this->BSE_A_local.size()); + if (is_full) assert(this->WB_local.size() == this->BSE_A_local.size()); + } +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::size_t i = 0; i < this->BSE_A_local.size(); ++i) { + if (alpha != 0.0) this->BSE_A_local[i] += alpha * this->VA_local[i]; + if (beta != 0.0) this->BSE_A_local[i] += beta * this->WA_local[i]; + if (is_full) { + if (alpha != 0.0) this->BSE_B_local[i] +=(alpha * this->VB_local[i]); + if (beta != 0.0) this->BSE_B_local[i] +=(beta * this->WB_local[i]); + } + } + } + // 3) check hermiticity/symmetry and (optionally) write to file + GlobalV::ofs_running << "CHECK hermiticity/symmetry" << std::endl; + std::cout << "| CHECK hermiticity/symmetry" << std::endl; + constexpr double threshold = 1.0e-6; + if (LR_Util::is_hermitian(this->BSE_A_local.data(), this->pA, threshold)) + { + std::cout << "| CHECK PASS: Matrix A is hermitian under threshold " << threshold << std::endl; + } + else + { + std::cout << "| CHECK WARNING: Matrix A is not hermitian under threshold " << threshold << std::endl; + } + if (PARAM.inp.out_bse_ab) + { + this->write_AB_matrix("A_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", 6, this->BSE_A_local.data(), this->ndim, this->ndim); + } + + if (is_full) + { + if (LR_Util::is_symmetric(this->BSE_B_local.data(), this->pA, threshold)) + { + std::cout << "| CHECK PASS: Matrix B is symmetric under threshold " << threshold << std::endl; + } + else + { + std::cout << "| CHECK WARNING: Matrix B is not symmetric under threshold " << threshold << std::endl; + } + if (PARAM.inp.out_bse_ab) + { + this->write_AB_matrix("B_matrix_"+std::to_string(GlobalV::MY_RANK)+".dat", 6, this->BSE_B_local.data(), this->ndim, this->ndim); + } + } + + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "init_bse_matrix"); + ModuleBase::timer::end("HamiltBSE", "init_bse_matrix"); +} + +template +void HamiltBSE::tda_solver(const int & st_index, const int& nstates, double* ene_out, T* X_out){ + ModuleBase::TITLE("HamiltBSE", "tda_solver"); + ModuleBase::timer::start("HamiltBSE", "elpa_tda_solver"); + + this->init_bse_matrix(false, st_index); + + std::vector X_tda(this->pA.get_local_size(), 0.0); + std::vector ev(nstates, 0.0); + + BSE::solve_tda(GlobalV::MY_RANK, + this->BSE_A_local, + this->pA, + ev, + X_tda); + // copy to output + std::copy_n(ev.data(), nstates, ene_out); + + LR_Util::pA2pX(X_out, X_tda.data(), nstates, this->nk, + this->nocc, this->nvirt, this->pX, this->pA, 0, 0, false/*openshell*/); + + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "BSE TDA solver"); + ModuleBase::timer::end("HamiltBSE", "elpa_tda_solver"); +} + +template <> +void HamiltBSE::full_solver(const int& st_index, const int& nstates, + double* ene_out, + double* X_out, + double* Y_out){ + ModuleBase::TITLE("HamiltBSE", "full_solver(double)"); + ModuleBase::timer::start("HamiltBSE", "elpa_full_solver(double)"); + + this->init_bse_matrix(true, st_index); + Parallel_2D pM; + LR_Util::setup_2d_division(pM, this->pA.get_block_size(), 2*this->ndim, 2*this->ndim + #ifdef __MPI + , this->pA.blacs_ctxt + #endif + ); + const auto to_complex = [](const std::vector& vin) { + return std::vector>(vin.begin(), vin.end()); + }; + std::vector> BSE_A_local_complex = to_complex(this->BSE_A_local); + std::vector> BSE_B_local_complex = to_complex(this->BSE_B_local); + std::vector> local_v_full(pM.get_local_size(), 0.0); + std::vector ev(2 * this->ndim, 0.0); + BSE::solve_full(GlobalV::MY_RANK, + BSE_A_local_complex, + BSE_B_local_complex, + this->pA, + pM, + ev, + local_v_full); + + std::vector local_v_full_real(pM.get_local_size(), 0.0); + for (size_t i = 0; i < local_v_full.size(); ++i) { + assert(std::abs(local_v_full[i].imag()) < 1e-10); + local_v_full_real[i] = local_v_full[i].real(); + } + + // copy positive eigenvalues to output + std::copy_n(&ev[this->ndim], nstates, ene_out); + LR_Util::pA2pX(X_out, local_v_full_real.data(), nstates, this->nk, + this->nocc, this->nvirt, this->pX, pM, 0, this->ndim, false/*openshell*/); + LR_Util::pA2pX(Y_out, local_v_full_real.data(), nstates, this->nk, + this->nocc, this->nvirt, this->pX, pM, this->ndim, this->ndim, false/*openshell*/); + + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "BSE Full solver"); + ModuleBase::timer::end("HamiltBSE", "elpa_full_solver(double)"); +} + +template <> +void HamiltBSE>::full_solver(const int& st_index, const int& nstates, + double* ene_out, + std::complex* X_out, + std::complex* Y_out){ + ModuleBase::TITLE("HamiltBSE", "full_solver(complex)"); + ModuleBase::timer::start("HamiltBSE", "elpa_full_solver(complex)"); + + this->init_bse_matrix(true, st_index); + Parallel_2D pM; + LR_Util::setup_2d_division(pM, this->pA.get_block_size(), 2*this->ndim, 2*this->ndim + #ifdef __MPI + , this->pA.blacs_ctxt + #endif + ); + + std::vector> local_v_full(pM.get_local_size(), 0.0); + std::vector ev(2 * this->ndim, 0.0); + ModuleBase::TITLE("HamiltBSE", "full_solver(complex)2"); + BSE::solve_full(GlobalV::MY_RANK, + this->BSE_A_local, + this->BSE_B_local, + this->pA, + pM, + ev, + local_v_full); + ModuleBase::TITLE("HamiltBSE", "full_solver(complex)3"); + // copy positive eigenvalues to output + std::copy_n(&ev[this->ndim], nstates, ene_out); + LR_Util::pA2pX(X_out, local_v_full.data(), nstates, this->nk, + this->nocc, this->nvirt, this->pX, pM, 0, this->ndim, false/*openshell*/); + LR_Util::pA2pX(Y_out, local_v_full.data(), nstates, this->nk, + this->nocc, this->nvirt, this->pX, pM, this->ndim, this->ndim, false/*openshell*/); + + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "BSE FULL solver"); + ModuleBase::timer::end("HamiltBSE", "elpa_full_solver(complex)"); +} + +template +void HamiltBSE::cal_V_by_grid(bool is_A) +{ + // 1. initialize HContainer VR + const int is = 0; //spin index, NOTE: only support 1 spin now + const auto psi_is = LR_Util::get_psi_spin(psi_ks, is, nk); + std::unique_ptr> VR = std::unique_ptr>(new hamilt::HContainer(&this->pmat)); + LR_Util::initialize_HR(*VR, this->ucell, this->gd, this->orb_cutoff); +#ifdef __MPI + Parallel_2D pV_col; //{nvirt*nocc, 1} + LR_Util::setup_2d_division(pV_col, this->pX[is].get_block_size(), nvirt[is]*nocc[is], 1, this->pX[is].blacs_ctxt); + std::vector V_col_local(this->nk * pV_col.get_local_size(), 0.0); // V_col(bjk2) +#endif + + int imo1, imo2; + for (int ik2 = 0; ik2 < nk; ++ik2) { + for (int j = 0; j < nocc[is]; ++j) { + for (int b = 0; b < nvirt[is]; ++b) {//calculate row {aik1} for each column {bjk2} + int bjk = ik2 * nocc[is] * nvirt[is] + j * nvirt[is] + b; // column index in BSE matrix + // 2. calculate transition matrix + if (is_A) //jk2→bk2, D(k)=c_b(k)c^†_j(k) + { imo1 = j; imo2 = b + nocc[is]; } + else //bjk2←kb2, D(k)=c_j(k)c^†_b(k) + { imo1 = b + nocc[is]; imo2 = j; } + #ifdef __MPI + ct::Tensor dm_trans_2d = + BSE_Util::cal_dm_trans_onebase_pblas(psi_is, pc, ik2, naos, imo1, imo2, pmat, (T)1.0 / (T)nk); + #else + ct::Tensor dm_trans_2d = + BSE_Util::cal_dm_trans_onebase_blas(psi_is, ik2, naos, imo1, imo2, (T)1.0 / (T)nk); + #endif + // LR_Util::print_tensor(dm_trans_2d, "dm_trans_2d", &pmat); + this->DM_trans->set_DMK_pointer(ik2, dm_trans_2d.data()); + // 3. D(k)→D(R) + this->DM_trans->cal_DMR(ik2); + // LR_Util::print_DMR(*DM_trans, ucell.nat, "DMR"); + + // 4. D(R)→V(R) + this->grid_calculation(*VR); + + // 5. V(R)→V(k) + std::vector v_k_2d(nk, LR_Util::newTensor({ pmat.get_col_size(), pmat.get_row_size() })); + for (auto& v : v_k_2d) v.zero(); + int nrow = ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) ? + this->pmat.get_row_size() : this->pmat.get_col_size(); + for (int ik1 = 0;ik1 < nk;++ik1) { + folding_HR(*VR, v_k_2d[ik1].data(), this->kv.kvec_d[ik1], nrow, 1); + } + // for (int ik1 = 0;ik1 < nk;++ik1) + // LR_Util::print_tensor(v_k_2d[ik1], "V(k)[ik=" + std::to_string(ik1) + "]", &this->pmat); + #ifdef __MPI + LR::ao_to_mo_pblas(v_k_2d, this->pmat, psi_is, this->pc, this->naos, + nocc[is], nvirt[is], pV_col, V_col_local.data(), false, LR_Util::MO_TYPE::VO); + + for (int ik1 = 0; ik1 < this->nk; ++ik1) { + Cpxgemr2d(nvirt[is]*nocc[is], 1, + V_col_local.data() + ik1 * pV_col.get_local_size(), 1, 1, pV_col.desc, + this->VA_local.data(), + ik1 * nocc[is] * nvirt[is] + 1 , bjk + 1, this->pA.desc, + this->pA.blacs_ctxt); + } + #else + LR::ao_to_mo_blas(v_k_2d, psi_is, nocc[is], nvirt[is], this->VA_local.data()+bjk * this->ndim, false, LR_Util::MO_TYPE::VO); + #endif + } + } + } +} + +template<> +void HamiltBSE::grid_calculation(hamilt::HContainer& VR) const +{ + ModuleBase::TITLE("HamiltBSE", "grid_calculation(double)"); + ModuleBase::timer::start("HamiltBSE", "grid_calculation(double)"); + + // 4.1. transition density rho on grid + double** rho_trans; + const int& nrxx = this->pot.lock()->nrxx; + + LR_Util::_allocate_2order_nested_ptr(rho_trans, 1, nrxx); // nspin=1 for transition density + ModuleBase::GlobalFunc::ZEROS(rho_trans[0], nrxx); + ModuleGint::cal_gint_rho(this->DM_trans->get_DMR_vector(), 1, rho_trans, false); + + // 4.2. v_hxc = f_hxc * rho_trans + ModuleBase::matrix vr_hxc(1, nrxx); //grid + std::vector ispin_ks = { 0 }; //for close-shell dft-xc kerenl, actually placeholder for bse + this->pot.lock()->cal_v_eff(rho_trans, ucell, vr_hxc, ispin_ks);// in this function, unit changes from Ha to Ry + LR_Util::_deallocate_2order_nested_ptr(rho_trans, 1); + + // 4.3 V^{Hxc}_{\mu,\nu}=\int{dr} \phi_\mu(r) v_{Hxc}(r) \phi_\nu(r) + VR.set_zero(); + ModuleGint::cal_gint_vl(vr_hxc.c, &VR); + // LR_Util::print_HR(VR, this->ucell.nat, "VR(real, 2d)"); + + ModuleBase::timer::end("HamiltBSE", "grid_calculation(double)"); +} + +template<> +void HamiltBSE>::grid_calculation(hamilt::HContainer>& VR) const +{ + ModuleBase::TITLE("HamiltBSE", "grid_calculation(complex)"); + ModuleBase::timer::start("HamiltBSE", "grid_calculation(complex)"); + + elecstate::DensityMatrix, double> DM_trans_real_imag(&this->pmat, 1, this->kv.kvec_d, this->nk); + DM_trans_real_imag.init_DMR(VR); + hamilt::HContainer HR_real_imag(ucell, &this->pmat); + LR_Util::initialize_HR, double>(HR_real_imag, ucell, gd, orb_cutoff); + + auto dmR_to_hR = [&, this](const char& type) -> void + { + LR_Util::get_DMR_real_imag_part(*this->DM_trans, DM_trans_real_imag, ucell.nat, type); + // if (this->first_print)LR_Util::print_DMR(DM_trans_real_imag, ucell.nat, "DMR(2d, real)"); + + // 4.1. transition density rho on grid + double** rho_trans; + const int& nrxx = this->pot.lock()->nrxx; + + LR_Util::_allocate_2order_nested_ptr(rho_trans, 1, nrxx); // nspin=1 for transition density + ModuleBase::GlobalFunc::ZEROS(rho_trans[0], nrxx); + ModuleGint::cal_gint_rho(DM_trans_real_imag.get_DMR_vector(), 1, rho_trans, false); + + // 4.2. v_hxc = f_hxc * rho_trans + ModuleBase::matrix vr_hxc(1, nrxx); //grid + std::vector ispin_ks = { 0 }; //for close-shell dft-xc kerenl, actually placeholder for bse + this->pot.lock()->cal_v_eff(rho_trans, ucell, vr_hxc, ispin_ks);// in this function, unit changes from Ha to Ry + LR_Util::_deallocate_2order_nested_ptr(rho_trans, 1); + + // 4.3 V^{Hxc}_{\mu,\nu}=\int{dr} \phi_\mu(r) v_{Hxc}(r) \phi_\nu(r) + HR_real_imag.set_zero(); + ModuleGint::cal_gint_vl(vr_hxc.c, &HR_real_imag); + // LR_Util::print_HR(HR_real_imag, this->ucell.nat, "VR(real, 2d)"); + LR_Util::set_HR_real_imag_part(HR_real_imag, VR, ucell.nat, type); + }; + VR.set_zero(); + dmR_to_hR('R'); //real + if (this->nk > 1) { dmR_to_hR('I'); } //imag for multi-k + ModuleBase::timer::end("HamiltBSE", "grid_calculation(complex)"); +} + +template class HamiltBSE; +template class HamiltBSE>; +}//namespace BSE \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/hamilt_bse.h b/source/source_lcao/module_lr/bse/hamilt_bse.h new file mode 100644 index 0000000000..4013a1cf49 --- /dev/null +++ b/source/source_lcao/module_lr/bse/hamilt_bse.h @@ -0,0 +1,120 @@ +// Purpose: simplify the hamilt_casida class, and explore to go beyond Tamm-Damcoff Approximation. +// For simplicity, only ELPA solver with MPI parallization is implementated. +// Thus, instead of iterative solver such as Davidson, here matrix is constructed directly. + +#pragma once +#include "bse_util.h" +#include "hamilt_bse_solver.h" +#include "molecular_lri.h" + +#include "source_base/parallel_2d.h" +#include "source_base/timer.h" +#include "source_basis/module_ao/parallel_orbitals.h" +#include "source_cell/unitcell.h" +#include "source_estate/module_dm/density_matrix.h" +#include "source_hamilt/hamilt.h" +#include "source_lcao/module_hcontainer/hcontainer_funcs.h" +#include "source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h" +#include "source_lcao/module_lr/potentials/pot_hxc_lrtd.h" +#include "source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h" +#include "source_lcao/module_lr/ri_benchmark/ri_benchmark.h" +#include "source_lcao/module_lr/utils/lr_io.h" +#include "source_lcao/module_lr/utils/lr_util.h" + +namespace BSE +{ +template +class HamiltBSE +{ +public: + std::vector BSE_A_local, BSE_B_local; + std::vector VA_local, WA_local; + std::vector VB_local, WB_local; + int ndim = 0; // dimension of BSE matrix + Parallel_2D pA; + /// @brief constructor for BSE_Matrix + HamiltBSE(const int& nspin, + const int& naos, + const std::vector& nocc, + const std::vector& nvirt, + const UnitCell& ucell_in, + const std::vector& orb_cutoff_in, + const Grid_Driver& gd_in, + const psi::Psi& psi_in, + const psi::Psi& psi_glb_in, + const ModuleBase::matrix& eig_gw_in, + MolecularLRI& mo_lri_in, + std::weak_ptr pot_in, + const K_Vectors& kv_in, + const std::vector& pX_in, + const Parallel_2D& pc_in, + const Parallel_Orbitals& pmat_in, + const std::vector& spin_types_in, + const std::string& tda, + const std::string& ri_hartree_benchmark_in = "none"); + /// @note these four functions shouldn't be called when `bse_mem_save` is true + void cal_V_for_A(); + void cal_W_for_A(); + void cal_V_for_B(); + void cal_W_for_B(); + + /// @brief initialize BSE matrix + /// @note if `bse_mem_save` is true, V and W matrix will be added directly + void init_bse_matrix(const bool is_full, const int& st_index); + + void tda_solver(const int& st_index, const int& nstates, double* ene_out, T* X_out); + void full_solver(const int& st_index, const int& nstates, double* ene_out, T* X_out, T* Y_out); + + void cal_V_by_grid(bool is_A); + void grid_calculation(hamilt::HContainer& VR) const; + + inline void write_AB_matrix(const std::string& file, const int& prec, const T* ptr, const int& size1, const int& size2) + { + std::ofstream ofs(file); + if (!ofs.is_open()){ + throw std::runtime_error("Cannot open file " + file); + } + ofs << file << "(Ry, transpose) with threshold " << prec << std::endl; + ofs << std::setprecision(prec) << std::scientific; + LR_Util::write_value(ofs, ptr, size1, size2); + ofs.close(); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "finish writing " + file); + } + + inline void read_AB_matrix(const std::string& file, T* ptr, const int& size1, const int& size2) + { + std::ifstream ifs(file); + if (!ifs.is_open()){ + throw std::runtime_error("Cannot open file " + file); + } + ifs.ignore(std::numeric_limits::max(), '\n'); // skip the first line + LR_Util::read_value(ifs, ptr, size1, size2); + ifs.close(); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "finish reading " + file); + } + +private: + const int nspin; + const int naos; + const std::vector nocc; + const std::vector nvirt; + const UnitCell& ucell; + const std::vector& orb_cutoff; + const Grid_Driver& gd; + const psi::Psi& psi_ks; + const psi::Psi& psi_ks_glb; + const ModuleBase::matrix& eig_gw; + MolecularLRI& mo_lri; + + std::weak_ptr pot; + const K_Vectors& kv; + int nk = 1; + const std::vector& pX; // for tda, also pY for full + const Parallel_2D& pc; + const Parallel_Orbitals& pmat; + const std::vector& spin_types; // singlet, triplet, and rpa, ipa(independent particle approx) + const std::string ri_hartree_benchmark; + + std::unique_ptr> DM_trans = nullptr; +}; +} // namespace BSE \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/hamilt_bse_solver.cpp b/source/source_lcao/module_lr/bse/hamilt_bse_solver.cpp new file mode 100644 index 0000000000..a2ae24b9d4 --- /dev/null +++ b/source/source_lcao/module_lr/bse/hamilt_bse_solver.cpp @@ -0,0 +1,331 @@ +#include "hamilt_bse_solver.h" + +namespace BSE +{ +template +void printM(const std::vector& A, int m, int n, std::string file, std::string name) +{ + std::ofstream ofs(file); + ofs << name << std::endl; + for (int i = 0; i < m; ++i) + { + for (int j = 0; j < n; ++j) + { + ofs << std::setw(25) << A[j * m + i]; + } + ofs << std::endl; + } + ofs.close(); +} + +void arrayFlatten1(const std::vector>& A, + const std::vector>& B, + std::vector>& result, // result={{A,B},{-B*,-A*}} + const Parallel_2D& pA, + const Parallel_2D& pM) +{ + int nA = pA.get_global_row_size(); +#ifdef __MPI + Cpxgemr2d(nA, nA, const_cast*>(A.data()), 1, 1, const_cast(pA.desc), + result.data(), 1, 1, const_cast(pM.desc), + pA.blacs_ctxt); + Cpxgemr2d(nA, nA, const_cast*>(B.data()), 1, 1, const_cast(pA.desc), + result.data(), 1, nA+1, const_cast(pM.desc), + pA.blacs_ctxt); + std::vector> temp_A(pA.get_local_size(), 0.0); + std::vector> temp_B(pA.get_local_size(), 0.0); + #pragma omp parallel for + for (int i = 0; i < pA.get_local_size(); i++ ){ + temp_A[i] = -std::conj(A[i]); + temp_B[i] = -std::conj(B[i]); + } + Cpxgemr2d(nA, nA, temp_A.data(), 1, 1, const_cast(pA.desc), + result.data(), nA+1, nA+1, const_cast(pM.desc), + pA.blacs_ctxt); + Cpxgemr2d(nA, nA, temp_B.data(), 1, 1, const_cast(pA.desc), + result.data(), nA+1, 1, const_cast(pM.desc), + pA.blacs_ctxt); +#else + for (int j = 0; j < nA; j++ ){ + for (int i = 0; i < nA; i++ ){ + result[ i + j*2*nA ] = A[i+j*nA]; + result[ nA+i + j*2*nA ] = -std::conj(B[i+j*nA]); + result[ i + (nA+j)*2*nA ] = B[i+j*nA]; + result[ nA+i + (nA+j)*2*nA ] = -std::conj(A[i+j*nA]); + } + } +#endif +} + +void arrayFlatten2(const std::vector>& A, + const std::vector>& B, + std::vector& result, // result={{Re(A+B), Im(A-B)},{-Im(A+B), Re(A-B)}} + const Parallel_2D& pA, + const Parallel_2D& pM) +{ + int nA = pA.get_global_row_size(); +#ifdef __MPI + std::vector temp(pA.get_local_size(), 0.0); + for (int i = 0; i < pA.get_local_size(); i++ ){ + temp[i] = A[i].real() + B[i].real(); + } + Cpxgemr2d(nA, nA, temp.data(), 1, 1, const_cast(pA.desc), + result.data(), 1, 1, const_cast(pM.desc), + pA.blacs_ctxt); + for (int i = 0; i < pA.get_local_size(); i++ ){ + temp[i] = A[i].imag() - B[i].imag(); + } + Cpxgemr2d(nA, nA, temp.data(), 1, 1, const_cast(pA.desc), + result.data(), 1, nA+1, const_cast(pM.desc), + pA.blacs_ctxt); + for (int i = 0; i < pA.get_local_size(); i++ ){ + temp[i] = - A[i].imag() - B[i].imag(); + } + Cpxgemr2d(nA, nA, temp.data(), 1, 1, const_cast(pA.desc), + result.data(), nA+1, 1, const_cast(pM.desc), + pA.blacs_ctxt); + for (int i = 0; i < pA.get_local_size(); i++ ){ + temp[i] = A[i].real() - B[i].real(); + } + Cpxgemr2d(nA, nA, temp.data(), 1, 1, const_cast(pA.desc), + result.data(), nA+1, nA+1, const_cast(pM.desc), + pA.blacs_ctxt); +#else +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int j = 0; j < nA; j++ ){ + for (int i = 0; i < nA; i++ ){ + result[ i + j*2*nA ] = A[i+j*nA].real() + B[i+j*nA].real(); + result[ nA+i + j*2*nA ] = - A[i+j*nA].imag() - B[i+j*nA].imag(); + result[ i + (nA+j)*2*nA ] = A[i+j*nA].imag() - B[i+j*nA].imag(); + result[ nA+i + (nA+j)*2*nA ] = A[i+j*nA].real() - B[i+j*nA].real(); + } + } +#endif +} + +void solve_full(const int my_rank, + std::vector>& A_part, + std::vector>& B_part, + const Parallel_2D& pA, + const Parallel_2D& pM, + std::vector& ev, + std::vector>& v) +{ + ModuleBase::TITLE("HamiltBSE", "elpa_solve_full"); + ModuleBase::timer::start("HamiltBSE", "elpa_solve_full"); + + assert(pM.get_global_row_size() == pM.get_global_col_size()); + int n = pM.get_global_row_size(); // full matrix size + int nA = pA.get_global_row_size(); + assert(n == 2*nA); + + elpa_t elpaInstance; + int status; + + if (elpa_init(20210430) != ELPA_OK) + { + fprintf(stderr, "Error: ELPA API version not supported"); + exit(1); + } + + elpaInstance = elpa_allocate(&status); + if (status != ELPA_OK) + { + std::cout << "Could not allocate elpa instance" << std::endl; + exit(1); + } + elpa_set(elpaInstance, "na", n, &status); + elpa_set(elpaInstance, "nev", n, &status); // have to solve all ev, see step6: v = SQLzΩ^(-1/2) + elpa_set(elpaInstance, "local_nrows", pM.get_row_size(), &status); + elpa_set(elpaInstance, "local_ncols", pM.get_col_size(), &status); + elpa_set(elpaInstance, "nblk", pM.get_block_size(), &status); + elpa_set(elpaInstance, "mpi_comm_parent", MPI_Comm_c2f(MPI_COMM_WORLD), &status); + elpa_set(elpaInstance, "process_row", pM.get_coord_row(), &status); + elpa_set(elpaInstance, "process_col", pM.get_coord_col(), &status); + elpa_set(elpaInstance, "solver", ELPA_SOLVER_2STAGE, &status); +#ifdef _OPENMP + int num_threads = omp_get_max_threads(); +#else + int num_threads = 1; +#endif + elpa_set(elpaInstance, "omp_threads", num_threads, &status); + status = elpa_setup(elpaInstance); + if (status != ELPA_OK) + { + fprintf(stderr, "Could not set up the ELPA object"); + } + + // step1: construct M + // M = {{Re(A_part+B_part), Im(A_part-B_part)}, {-Im(A_part+B_part), Re(A_part-B_part)}} + ModuleBase::TITLE("HamiltBSE", "elpa_solve_full1"); + std::vector M(pM.get_local_size()); + arrayFlatten2(A_part, B_part, M, pA, pM); + std::vector>().swap(A_part); + std::vector>().swap(B_part); + + // step2: construct J = {{0, I}, {-I, 0}} + ModuleBase::TITLE("HamiltBSE", "elpa_solve_full2"); + std::vector J(pM.get_local_size(), 0.0); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int j = 0;j < pM.get_col_size();++j) + { + int col_glb = pM.local2global_col(j); + for (int i = 0;i < pM.get_row_size();++i) + { + int row_glb = pM.local2global_row(i); + if (col_glb - row_glb == nA) + J[j * pM.get_row_size() + i] = 1.0; + else if (row_glb - col_glb == nA) + J[j * pM.get_row_size() + i] = -1.0; + } + } + + // step3: Cholesky factorization M = U^T U + ModuleBase::TITLE("HamiltBSE", "elpa_solve_full3"); + elpa_cholesky(elpaInstance, M.data(), &status); // output M is upper triangular matrix U, M=U^T U + + // std::vector global_U(n*n, 0.0); + // Cpxgemr2d(n, n, M.data(), 1, 1, const_cast(pM.desc), + // global_U.data(), 1, 1, pM_glb.desc, + // pM.blacs_ctxt); + // if (my_rank == 0) {std::cout<<"U:"< UJ (pM.get_local_size(), 0.0); + + ScalapackConnector::gemm('N', 'N', n, n, n, 1.0, + M.data(), 1, 1, pM.desc, + J.data(), 1, 1, pM.desc, + 0.0, + UJ.data(), 1, 1, pM.desc); + + // std::vector global_UJ (n*n); + // Cpxgemr2d(n, n, UJ.data(), 1, 1, const_cast(pM.desc), + // global_UJ.data(), 1, 1, pM_glb.desc, + // pM.blacs_ctxt); + //if (my_rank == 0) {std::cout<<"UJ:"<().swap(UJ); + + // std::vector global_UJL(n*n); + // Cpxgemr2d(n, n, J.data(), 1, 1, const_cast(pM.desc), + // global_UJL.data(), 1, 1, pM_glb.desc, + // pM.blacs_ctxt); + //if (my_rank == 0) {std::cout<<"UJL:"< z(2 * pM.get_local_size()); // 2 for elpa_skew stores complex as 2 double + + elpa_skew_eigenvectors(elpaInstance, J.data()/*UJL*/, ev.data(), z.data(), &status); + assert(status == ELPA_OK); + std::vector().swap(J); + // std::vector> global_z(n * n, 0.0); + // Cpxgemr2d(n, n, z.data(), 1, 1, const_cast(pM.desc), + // global_z.data(), 1, 1, pM_glb.desc, + // pM.blacs_ctxt); + + // if (my_rank == 0) { + // std::cout << "Eigenvalues computed successfully." << std::endl; + // std::cout << "Eigenvalues:" << std::endl; + // for (int i = 0; i < n; ++i) { + // std::cout << ev[i] << std::endl; + // } + //std::cout << "Eigenvectors (right):" << std::endl; + //printM(global_z, n, n); + //} + // step6: compute normalized eigenvectors v = SQLzΩ^(-1/2), where Ω = diag(ev) + ModuleBase::TITLE("HamiltBSE", "elpa_solve_full6.1"); + std::vector Lz_real(pM.get_local_size(), 0.0); + std::vector Lz_imag(pM.get_local_size(), 0.0); + +// 6.1: zΩ^(-1/2). NOTE: both positive and negative eigenvalues are handled here +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int j = 0; j < pM.get_col_size(); ++j) + { + double ev_sqrt = std::sqrt(std::abs(ev[pM.local2global_col(j)])); + for (int i = 0; i < pM.get_row_size(); ++i) + { + z[j * pM.get_row_size() + i] /= ev_sqrt; // real part + z[j * pM.get_row_size() + i + pM.get_local_size()] /= ev_sqrt; // imaginary part + } + } + MPI_Barrier(MPI_COMM_WORLD); + + // 6.2: LzΩ^(-1/2), combine real and imaginary part + ModuleBase::TITLE("HamiltBSE", "elpa_solve_full6.2"); + ScalapackConnector::gemm('T', 'N', n, n, n, 1.0, + M.data(), 1, 1, pM.desc, + z.data(), 1, 1, pM.desc, + 0.0, + Lz_real.data(), 1, 1, pM.desc); + ScalapackConnector::gemm('T', 'N', n, n, n, 1.0, + M.data(), 1, 1, pM.desc, + z.data() + pM.get_local_size(), 1, 1, pM.desc, + 0.0, + Lz_imag.data(), 1, 1, pM.desc); + + std::vector> Lz(pM.get_local_size(), 0.0); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int j = 0; j < pM.get_col_size(); ++j) + { + for (int i = 0; i < pM.get_row_size(); ++i) + { + Lz[j * pM.get_row_size() + i] + = std::complex(Lz_real[j * pM.get_row_size() + i], Lz_imag[j * pM.get_row_size() + i]); + } + } + std::vector().swap(Lz_real); + std::vector().swap(Lz_imag); + + std::vector> SQ(pM.get_local_size(), 0.0);// SQ = {{I, -i I}, {-I, -i I}} / sqrt(2) +#ifdef _OPENMP +#pragma omp parallel for +#endif + for (int j = 0; j < pM.get_col_size(); ++j) { + int col_glb = pM.local2global_col(j); + for (int i = 0; i < pM.get_row_size(); ++i) { + int row_glb = pM.local2global_row(i); + if ( (col_glb < nA) && (row_glb < nA) && (col_glb == row_glb) ) { + SQ[j * pM.get_row_size() + i] = 1/std::sqrt(2); + } + else if ( (col_glb >= nA) && (row_glb >= nA) && (col_glb == row_glb) ) { + SQ[j * pM.get_row_size() + i] = {0, -1/std::sqrt(2)}; + } + else if (col_glb == row_glb - nA) { + SQ[j * pM.get_row_size() + i] = -1/std::sqrt(2); + } + else if (col_glb - nA == row_glb) { + SQ[j * pM.get_row_size() + i] = {0, -1/std::sqrt(2)}; + } + } + } + MPI_Barrier(MPI_COMM_WORLD); + // 6.3: v=SQLzΩ^{-1/2} + ModuleBase::TITLE("HamiltBSE", "elpa_solve_full6.3"); + ScalapackConnector::gemm('N', 'N', n, n, n, 1.0, + SQ.data(), 1, 1, pM.desc, + Lz.data(), 1, 1, pM.desc, + 0.0, + v.data(), 1, 1, pM.desc); + + elpa_deallocate(elpaInstance, &status); + elpa_uninit(&status); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "elpa_solve_full"); + ModuleBase::timer::end("HamiltBSE", "elpa_solve_full"); +} +} // namespace BSE \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/hamilt_bse_solver.h b/source/source_lcao/module_lr/bse/hamilt_bse_solver.h new file mode 100644 index 0000000000..8dc5fea4ed --- /dev/null +++ b/source/source_lcao/module_lr/bse/hamilt_bse_solver.h @@ -0,0 +1,53 @@ +#pragma once + +#include "source_base/parallel_2d.h" +#include "source_base/parallel_reduce.h" +#include "source_base/module_external/blacs_connector.h" +#include "source_base/module_external/scalapack_connector.h" +#include "source_lcao/module_lr/utils/lr_util.h" +#define HAVE_SKEWSYMMETRIC // for elpa_skew_eigenvectors +#include +#ifdef I // to avoid conflict with macro I defined in elpa +#undef I +#endif +#include + +namespace BSE +{ +template +void printM(const std::vector& A, int m, int n, std::string file, std::string name); + +/// @brief result={{A,B},{-A*,-B*}} +void arrayFlatten1(const std::vector>& A, + const std::vector>& B, + std::vector>& result, + const Parallel_2D& pA, + const Parallel_2D& pM); + +/// @brief result={{Re(A+B), Im(A-B)},{-Im(A+B), Re(A-B)}} +void arrayFlatten2(const std::vector>& A, + const std::vector>& B, + std::vector& result, + const Parallel_2D& pA, + const Parallel_2D& pM); + +/// @brief solve full BSE through matrix M = {{Re(A+B), Im(A-B)},{-Im(A+B), Re(A-B)}} +/// @attention A_part and B_part will be destroyed (swapped with empty vectors) +/// after constructing M. Pass copies if the caller needs them afterwards. +void solve_full(const int my_rank, + std::vector>& A_part, + std::vector>& B_part, + const Parallel_2D& pA, + const Parallel_2D& pM, + std::vector& ev, + std::vector>& global_v); +/// @brief template solve TDA BSE, implemented in hamilt_bse_solver.hpp +template +void solve_tda(const int my_rank, + const std::vector& A, + const Parallel_2D& pA, + std::vector& ev, + std::vector& v); +} + +#include "hamilt_bse_solver.hpp" \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/hamilt_bse_solver.hpp b/source/source_lcao/module_lr/bse/hamilt_bse_solver.hpp new file mode 100644 index 0000000000..6300ee1369 --- /dev/null +++ b/source/source_lcao/module_lr/bse/hamilt_bse_solver.hpp @@ -0,0 +1,62 @@ +#pragma once + +#include "hamilt_bse_solver.h" +namespace BSE +{ +template +void solve_tda(const int my_rank, + const std::vector& A, + const Parallel_2D& pA, + std::vector& ev, + std::vector& v) +{ + ModuleBase::TITLE("HamiltBSE", "elpa_solve_tda1"); + ModuleBase::timer::start("HamiltBSE", "elpa_solve_tda"); + + assert(pA.get_global_row_size() == pA.get_global_col_size()); + const int nA = pA.get_global_row_size(); + + elpa_t elpaInstance; + int status; + + if (elpa_init(20210430) != ELPA_OK) + { + fprintf(stderr, "Error: ELPA API version not supported"); + exit(1); + } + + elpaInstance = elpa_allocate(&status); + if (status != ELPA_OK) + { + std::cout << "Could not allocate elpa instance" << std::endl; + exit(1); + } + elpa_set(elpaInstance, "na", nA, &status); + elpa_set(elpaInstance, "nev", static_cast(ev.size()), &status); + elpa_set(elpaInstance, "local_nrows", pA.get_row_size(), &status); + elpa_set(elpaInstance, "local_ncols", pA.get_col_size(), &status); + elpa_set(elpaInstance, "nblk", pA.get_block_size(), &status); + elpa_set(elpaInstance, "mpi_comm_parent", MPI_Comm_c2f(MPI_COMM_WORLD), &status); + elpa_set(elpaInstance, "process_row", pA.get_coord_row(), &status); + elpa_set(elpaInstance, "process_col", pA.get_coord_col(), &status); + elpa_set(elpaInstance, "solver", ELPA_SOLVER_2STAGE, &status); +#ifdef _OPENMP + int num_threads = omp_get_max_threads(); +#else + int num_threads = 1; +#endif + elpa_set(elpaInstance, "omp_threads", num_threads, &status); + status = elpa_setup(elpaInstance); + if (status != ELPA_OK) + { + fprintf(stderr, "Could not set up the ELPA object"); + } + ModuleBase::TITLE("HamiltBSE", "elpa_solve_tda2"); + elpa_eigenvectors(elpaInstance, const_cast(A.data()), ev.data(), v.data(), &status); + ModuleBase::TITLE("HamiltBSE", "elpa_solve_tda3"); + elpa_deallocate(elpaInstance, &status); + elpa_uninit(&status); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "elpa_solve_tda"); + ModuleBase::timer::end("HamiltBSE", "elpa_solve_tda"); +} +} \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/molecular_lri.h b/source/source_lcao/module_lr/bse/molecular_lri.h new file mode 100644 index 0000000000..bc1ac70c82 --- /dev/null +++ b/source/source_lcao/module_lr/bse/molecular_lri.h @@ -0,0 +1,164 @@ +//======================= +// AUTHOR : Ziqing Guan +// DATE : 2026-03-22 +//======================= +#pragma once +#include +#include "source_base/timer.h" +#include "source_base/global_function.h" +#include "source_base/module_container/base/third_party/blas.h" +#include "source_cell/unitcell.h" +#include "source_cell/klist.h" +#include "source_hamilt/module_xc/exx_info_ri.h" +#include "source_lcao/module_ri/LRI_CV_Tools.h" +#include "source_lcao/module_lr/bse/bse_util.h" +#include "source_lcao/module_lr/utils/lr_util.h" +#include "source_lcao/module_lr/utils/lr_util_print.h" +#include "source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h" +#include "source_lcao/module_lr/utils/lr_io.h" +namespace BSE +{ + +using TA = int; +using TC = std::array; +using Tk = std::array; +using TAC = std::pair; +using TatomR = std::array; + +template +using TLRI = std::map>>; + +template +class MolecularLRI +{ +public: + RI::LR LR_lri; + RI::Cell_Nearest cell_nearest; + /// @brief calculate V[k_ai][k_bj](j,b,i,a) and W[k_ai][k_bj](j,b,i,a) with RI method + MolecularLRI(const UnitCell& ucell, + const int nk, + const LR_IO::RI_kRlist& kRlist_in, + const int nocc, + const int nvirt, + const psi::Psi& psi_ks_in) // < ATTENTION: psi_ks should be global + : ucell(ucell), nk(nk), kRlist(kRlist_in), kv(*kRlist_in.klist), nocc(nocc), nvirt(nvirt), + ndim(nk*nocc*nvirt), psi_ks(psi_ks_in), is_local_k1(nk, false) + { + std::map atoms_pos; + for (int iat = 0; iat < this->ucell.nat; ++iat) + { + atoms_pos[iat] = RI_Util::Vector3_to_array3( + this->ucell.atoms[this->ucell.iat2it[iat]].tau[this->ucell.iat2ia[iat]]); + } + const std::array latvec = {RI_Util::Vector3_to_array3(this->ucell.a1), + RI_Util::Vector3_to_array3(this->ucell.a2), + RI_Util::Vector3_to_array3(this->ucell.a3)}; + this->LR_lri.set_parallel(MPI_COMM_WORLD, atoms_pos, latvec, kRlist.period); + this->cell_nearest.init(atoms_pos, latvec, kRlist.period); + }; + + ~MolecularLRI() {} + + /// =============== calculation interface ==================== + void init(TLRI& Cs_in, TLRI& Vs_in, TLRI& Ws_in, const Exx_Info_RI& info_ri); + + void cal_W_for_A(std::vector& m_2d, const Parallel_2D& pm_2d, const double factor=1.0) + { + ModuleBase::TITLE("MolecularLRI", "cal_W_for_A"); + ModuleBase::timer::start("MolecularLRI", "cal_W_for_A"); + std::map>> + Wk = this->LR_lri.cal_cvc_mo_k_onthefly({"O","O","V","V"}, "Ws_", true); + ModuleBase::timer::end("MolecularLRI", "cal_W_for_A"); + this->transform_k_2dlocal(m_2d, Wk, pm_2d, factor); + } + void cal_W_for_B(std::vector& m_2d, const Parallel_2D& pm_2d, const double factor=1.0) + { + ModuleBase::TITLE("MolecularLRI", "cal_W_for_B"); + ModuleBase::timer::start("MolecularLRI", "cal_W_for_B"); + std::map>> + Wk = this->LR_lri.cal_cvc_mo_k_onthefly({"V","O","O","V"}, "Ws_", false); + ModuleBase::timer::end("MolecularLRI", "cal_W_for_B"); + this->transform_k_2dlocal(m_2d, Wk, pm_2d, factor); + } + void cal_hartree_for_A(std::vector& m_2d, const Parallel_2D& pm_2d, const double factor=1.0) + { + ModuleBase::TITLE("MolecularLRI", "cal_hartree_for_A"); + ModuleBase::timer::start("MolecularLRI", "cal_hartree_for_A"); + std::map>> + Vk = this->LR_lri.cal_cvc_mo_k_hartree_onthefly({"O","V","O","V"}, "Vs_", true); + ModuleBase::timer::end("MolecularLRI", "cal_hartree_for_A"); + this->transform_k_2dlocal(m_2d, Vk, pm_2d, factor); + } + void cal_hartree_for_B(std::vector& m_2d, const Parallel_2D& pm_2d, const double factor=1.0) + { + ModuleBase::TITLE("MolecularLRI", "cal_hartree_for_B"); + ModuleBase::timer::start("MolecularLRI", "cal_hartree_for_B"); + std::map>> + Vk = this->LR_lri.cal_cvc_mo_k_hartree_onthefly({"O","V","O","V"}, "Vs_", false); + ModuleBase::timer::end("MolecularLRI", "cal_hartree_for_B"); + this->transform_k_2dlocal(m_2d, Vk, pm_2d, factor); + } + + /// =============== print ==================== + inline void print_k(std::ostream& ofs, const std::vector& kindex_map, const std::vector& vec, const std::string name) + { + ofs << name << ": size = " << vec.size() << std::endl; + ofs << std::fixed << std::setprecision(4); + int count = 0; + for (auto v : vec) + { + Tk k = kindex_map[v]; + ofs << "(" << std::setw(6) << k[0] <<", "<< std::setw(6) << k[1] <<", "<< std::setw(6) << k[2] <<") "; + count++; + if (count % 5 == 0) { ofs << std::endl; } + } + ofs << std::endl; + ofs << std::defaultfloat; + } + inline void print_a(std::ostream& ofs, const std::vector& vec, const std::string name) + { + ofs << name << ": "; + int count = 0; + for (auto& v : vec) + { + ofs << v << " "; + count++; + if (count % 10 == 0) { ofs << std::endl; } + } + ofs << std::endl; + } + +protected: + /// =============== inner function ==================== + void transform_k_2dlocal(std::vector& m_2d, + const std::map>>& m_lri, + const Parallel_2D& pm_2d, const double factor); + + // > + std::map>> cal_Csk_ao_mo(const TLRI& CsR_ao, + const std::vector& k_list, + const std::vector& list_IJ); + + // transform total psi to map type according k coordinate and atom index + std::map>> transform_psi_k(const psi::Psi& psi_ks, + const std::vector& k_list); + + void build_q_to_kpair_map(int mode, double threshold); + + const UnitCell& ucell; + const int nk; + const K_Vectors& kv; + const LR_IO::RI_kRlist& kRlist; + const Parallel_2D pm_2d; + const int nocc; + const int nvirt; + const int ndim; + const psi::Psi& psi_ks; + std::vector is_local_k1; + +};// class MolecularLRI + +}// namespace BSE + +#include "molecular_lri.hpp" +#include "molecular_lri_comm.hpp" diff --git a/source/source_lcao/module_lr/bse/molecular_lri.hpp b/source/source_lcao/module_lr/bse/molecular_lri.hpp new file mode 100644 index 0000000000..5490f70d70 --- /dev/null +++ b/source/source_lcao/module_lr/bse/molecular_lri.hpp @@ -0,0 +1,289 @@ +//======================= +// AUTHOR : Ziqing Guan +// DATE : 2026-03-22 +//======================= +#include "molecular_lri.h" +#include +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif + +#ifdef __MKL +#include +#endif +namespace BSE +{ + +template +void MolecularLRI::init(TLRI& Cs_in, TLRI& Vs_in, TLRI& Ws_in, const Exx_Info_RI& info_ri) +{ + ModuleBase::TITLE("MolecularLRI", "init"); + ModuleBase::timer::start("MolecularLRI", "init"); + // 1. distribute atom and k-point tasks among MPI processes + int nproc = GlobalV::NPROC; + std::set set_I, set_J, set_IJ, set_k; + int task_sizes = this->ucell.nat * this->ucell.nat * this->nk * this->nk; + std::cout << "Total Molecular LRI tasks: " << task_sizes << ", number of MPI processes: " << nproc << std::endl; + RI::Distribute_Equally::distribute_atom_and_k_pair(MPI_COMM_WORLD, + (std::size_t)this->ucell.nat, + (std::size_t)this->nk, + this->LR_lri.list_I, + this->LR_lri.list_J, + this->LR_lri.k1_indices, + this->LR_lri.k2_indices, + false); + + set_IJ.insert(this->LR_lri.list_I.begin(), this->LR_lri.list_I.end()); + set_IJ.insert(this->LR_lri.list_J.begin(), this->LR_lri.list_J.end()); + // transform set to vector, since openmp cannot handle set for parallelization + this->LR_lri.list_IJ.assign(set_IJ.begin(), set_IJ.end()); + + for (int k1 : this->LR_lri.k1_indices) + { + this->is_local_k1[k1] = true; + } + set_k.insert(this->LR_lri.k1_indices.begin(), this->LR_lri.k1_indices.end()); + set_k.insert(this->LR_lri.k2_indices.begin(), this->LR_lri.k2_indices.end()); + this->LR_lri.k_indices.assign(set_k.begin(), set_k.end()); + + // build kindex_map: index → fractional coordinate for all k-points + this->LR_lri.kindex_map.resize(this->nk); + for (int ik = 0; ik < this->nk; ++ik) + this->LR_lri.kindex_map[ik] = RI_Util::Vector3_to_array3(this->kv.kvec_d.at(ik)); + + int proc_ntasks = this->LR_lri.list_I.size() * this->LR_lri.list_J.size() * this->LR_lri.k1_indices.size() * this->LR_lri.k2_indices.size(); + GlobalV::ofs_running << "Molecular LRI init: Process " << GlobalV::MY_RANK + << " handles " << proc_ntasks << " tasks." << std::endl; + print_a(GlobalV::ofs_running, this->LR_lri.list_I, "list_I"); + print_a(GlobalV::ofs_running, this->LR_lri.list_J, "list_J"); + print_k(GlobalV::ofs_running, this->LR_lri.kindex_map, this->LR_lri.k1_indices, "list_k1"); + print_k(GlobalV::ofs_running, this->LR_lri.kindex_map, this->LR_lri.k2_indices, "list_k2"); + // LRI_CV_Tools::write_Cs_ao(Cs_in, PARAM.globalv.global_out_dir + "Cs_in_test_" + std::to_string(GlobalV::MY_RANK)); + // LRI_CV_Tools::write_Vs_abf(Vs_in, PARAM.globalv.global_out_dir + "Vs_in_test_" + std::to_string(GlobalV::MY_RANK)); + // LRI_CV_Tools::write_Vs_abf(Ws_in, PARAM.globalv.global_out_dir + "Ws_in_test_" + std::to_string(GlobalV::MY_RANK)); + + // build q→kpair mapping (for W term q-approximation) + this->build_q_to_kpair_map(PARAM.inp.bse_q_approx_mode, PARAM.inp.bse_q_approx_threshold); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "distribute_atom_and_k"); + + // 2. move R of tensors to nearest image Bvk cell for valid Fourier interpolation + double dist; + std::set all_atoms; + for (int i = 0; i < this->ucell.nat; ++i) + { + all_atoms.insert(i); + for (int j = 0; j < this->ucell.nat; ++j) + { + for (const TC R_original : this->kRlist.Rlist) + { + const TC R = cell_nearest.cell_nearest_direction(i, j, R_original, dist); + if (R != R_original) + { + BSE_Util::move_R_tensor(Cs_in, i, j, R_original, R); + BSE_Util::move_R_tensor(Vs_in, i, j, R_original, R); + BSE_Util::move_R_tensor(Ws_in, i, j, R_original, R); + } + } + } + } + + + // 3. set tensors, in these functions MPI distribution will be performed + ModuleBase::TITLE("MolecularLRI", "before_set_Cs"); + this->LR_lri.set_Cs(Cs_in, info_ri.C_threshold, set_IJ, all_atoms); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "set_Cs"); + ModuleBase::TITLE("MolecularLRI", "before_set_Vs"); + this->LR_lri.set_Vs(Vs_in, info_ri.V_threshold, set_I, set_J); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "set_Vs"); + ModuleBase::TITLE("MolecularLRI", "before_set_Ws"); + this->LR_lri.set_Ws(Ws_in, info_ri.V_threshold, set_I, set_J); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "set_Ws"); + if (PARAM.inp.out_ri_cv) + { + TLRI& Cs_LRI = this->LR_lri.lri.data_pool.at("Cs_").Ds_ab;// see LRI::set_tensor_map2 + LRI_CV_Tools::write_Cs_ao(Cs_LRI, PARAM.globalv.global_out_dir + "Cs_lri_test_" + std::to_string(GlobalV::MY_RANK)); + TLRI& Vs_LRI = this->LR_lri.lri.data_pool.at("Vs_").Ds_ab; + LRI_CV_Tools::write_Vs_abf(Vs_LRI, PARAM.globalv.global_out_dir + "Vs_lri_test_" + std::to_string(GlobalV::MY_RANK)); + TLRI& Ws_LRI = this->LR_lri.lri.data_pool.at("Ws_").Ds_ab; + LRI_CV_Tools::write_Vs_abf(Ws_LRI, PARAM.globalv.global_out_dir + "Ws_lri_test_" + std::to_string(GlobalV::MY_RANK)); + } + // 4. prepare mo-type tensors + this->LR_lri.nocc = this->nocc; + this->LR_lri.nvirt = this->nvirt; + this->LR_lri.map_psi = this->transform_psi_k(this->psi_ks, this->LR_lri.k_indices); + + ModuleBase::TITLE("MolecularLRI", "cal_Csk_ao_mo"); + ModuleBase::timer::start("MolecularLRI", "cal_Csk_ao_mo"); + this->LR_lri.cal_Csk_ao_mo("Cs_", GlobalV::ofs_running); + ModuleBase::timer::end("MolecularLRI", "cal_Csk_ao_mo"); + + ModuleBase::TITLE("MolecularLRI", "before_free_Cs"); + this->LR_lri.free_Cs(); // free Cs_ao to save memory + ModuleBase::TITLE("MolecularLRI", "after_free_Cs"); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "MolecularLRI_init"); + ModuleBase::timer::end("MolecularLRI", "init"); +} + +/// @brief build q_list and q2kpair mapping according to approximation mode +template +void MolecularLRI::build_q_to_kpair_map(int mode, double threshold) +{ + ModuleBase::TITLE("MolecularLRI", "build_q_to_kpair_map"); + using namespace RI::Array_Operator; + const Tk k_unit{1.0, 1.0, 1.0}; + + // Fuzzy comparator for q_set + struct kComparator { + bool operator()(const Tk& lhs, const Tk& rhs) const { + constexpr double eps = 1e-6; + for (int i = 0; i < 3; ++i) + if (std::abs(lhs[i] - rhs[i]) > eps) return lhs[i] < rhs[i]; + return false; + } + }; + + auto frac_to_cart = [&](const Tk& k_frac) -> Tk { + ModuleBase::Vector3 v_frac(k_frac[0], k_frac[1], k_frac[2]); + ModuleBase::Vector3 v_cart = v_frac * ucell.G; + return Tk{v_cart.x, v_cart.y, v_cart.z}; + }; + auto cart_to_frac = [&](const Tk& k_cart) -> Tk { + ModuleBase::Vector3 v_cart(k_cart[0], k_cart[1], k_cart[2]); + ModuleBase::Vector3 v_frac = v_cart * ucell.latvec.Transpose(); + return Tk{v_frac.x, v_frac.y, v_frac.z}; + }; + auto cart_dist = [&](const Tk& q1, const Tk& q2) -> double { + Tk dq_frac = cart_to_frac({q1[0]-q2[0], q1[1]-q2[1], q1[2]-q2[2]}); + for (int i = 0; i < 3; ++i) dq_frac[i] -= std::round(dq_frac[i]); + Tk dq_cart = frac_to_cart(dq_frac); + return std::sqrt(dq_cart[0]*dq_cart[0] + dq_cart[1]*dq_cart[1] + dq_cart[2]*dq_cart[2]); + }; + + if (mode == 0) + { + std::map>, kComparator> q2kpair_tmp; + for (int k1 : this->LR_lri.k1_indices) + for (int k2 : this->LR_lri.k2_indices) + { + Tk q = (this->LR_lri.kindex_map[k2] - this->LR_lri.kindex_map[k1]) % k_unit; + q2kpair_tmp[q].emplace_back(k1, k2); + } + this->LR_lri.q2kpair.insert(std::make_move_iterator(q2kpair_tmp.begin()), std::make_move_iterator(q2kpair_tmp.end())); + } + else + { + std::set q_coarse_set; + const K_Vectors& kv_coarse = this->kRlist.klist_coarse; + int nk_coarse = kv_coarse.get_nkstot_full(); + for (int ik1 = 0; ik1 < nk_coarse; ++ik1) + { + Tk ck1 = RI_Util::Vector3_to_array3(kv_coarse.kvec_d.at(ik1)); + for (int ik2 = 0; ik2 < nk_coarse; ++ik2) + { + Tk ck2 = RI_Util::Vector3_to_array3(kv_coarse.kvec_d.at(ik2)); + q_coarse_set.insert((ck2 - ck1) % k_unit); + } + } + std::vector q_coarse_list(q_coarse_set.begin(), q_coarse_set.end()); + std::vector q_coarse_cart; + for (const Tk& q : q_coarse_list) q_coarse_cart.push_back(frac_to_cart(q)); + + for (int k1 : this->LR_lri.k1_indices) + { + Tk k1_frac = this->LR_lri.kindex_map[k1]; + for (int k2 : this->LR_lri.k2_indices) + { + Tk k2_frac = this->LR_lri.kindex_map[k2]; + Tk q_frac = (k2_frac - k1_frac) % k_unit; + Tk q_cart = frac_to_cart(q_frac); + double q_norm = std::sqrt(q_cart[0]*q_cart[0] + q_cart[1]*q_cart[1] + q_cart[2]*q_cart[2]); + + Tk q_target; + if (mode == 2 && q_norm < threshold) + { + q_target = q_frac; + } + else + { + double min_dist = std::numeric_limits::max(); + int nearest = -1; + for (std::size_t iq = 0; iq < q_coarse_cart.size(); ++iq) + { + double d = cart_dist(q_cart, q_coarse_cart[iq]); + if (d < min_dist) { min_dist = d; nearest = iq; } + } + q_target = q_coarse_list[nearest]; + } + this->LR_lri.q2kpair[q_target].emplace_back(k1, k2); + } + } + } + + for (auto& pair : this->LR_lri.q2kpair) + { + this->LR_lri.q_list.push_back(pair.first); + } + // Print q_list with both direct and Cartesian coordinates + GlobalV::ofs_running << "q_list: size = " << this->LR_lri.q_list.size() << std::endl; + GlobalV::ofs_running << "detailed q_list see 'qlist_{rank}.dat'" << std::endl; + std::ofstream ofs_qlist(PARAM.globalv.global_out_dir + "qlist_" + std::to_string(GlobalV::MY_RANK) + ".dat"); + ofs_qlist << "q_list: size = " << this->LR_lri.q_list.size() << std::endl; + ofs_qlist << "(direct coords) | (Cartesian, Bohr^-1)" << std::endl; + ofs_qlist << std::fixed << std::setprecision(4); + for (const Tk& q : this->LR_lri.q_list) + { + Tk qc = frac_to_cart(q); + ofs_qlist << "(" << std::setw(6) << q[0] << "," << std::setw(6) << q[1] + << "," << std::setw(6) << q[2] << ") | (" << std::setw(7) << qc[0] + << "," << std::setw(7) << qc[1] << "," << std::setw(7) << qc[2] << ")" << std::endl; + } + ofs_qlist << std::defaultfloat; +} + +/// @brief transform psi to >, mo is not sliced +template +std::map>> +MolecularLRI::transform_psi_k(const psi::Psi& psi_ks, const std::vector& k_list) +{ + ModuleBase::TITLE("MolecularLRI", "transform_psi_k"); + ModuleBase::timer::start("MolecularLRI", "transform_psi_k"); + std::map>> psi_map; + const std::size_t nmo = psi_ks.get_nbands(); + for (const int ik : k_list) // initialize + { + auto& psi_map_k = psi_map[ik]; + for (int iat = 0; iat < this->ucell.nat; ++iat) + { + psi_map_k[iat]; + } + } +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (const auto& ik : k_list) + { + auto& psi_map_k = psi_map.at(ik); + for (int iat = 0; iat < this->ucell.nat; ++iat) + { + const int it = this->ucell.iat2it[iat]; + const std::size_t nw = this->ucell.atoms[it].nw; + RI::Tensor t({nmo, nw}); + for (int im = 0; im < nmo; ++im) + { + for (int iw = 0; iw < nw; ++iw) + { + t(im, iw) = psi_ks(ik, im, this->ucell.get_iat2iwt()[iat]+iw); + } + } + psi_map_k.at(iat) = std::move(t); + } + } + ModuleBase::timer::end("MolecularLRI", "transform_psi_k"); + return psi_map; +} + +}// namespace BSE \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/molecular_lri_comm.hpp b/source/source_lcao/module_lr/bse/molecular_lri_comm.hpp new file mode 100644 index 0000000000..5dd9f8d45a --- /dev/null +++ b/source/source_lcao/module_lr/bse/molecular_lri_comm.hpp @@ -0,0 +1,276 @@ +//======================= +// AUTHOR : Ziqing Guan +// DATE : 2026-03-22 +//======================= +#include "molecular_lri.h" +#include +#include // offsetof +#include +#ifdef _OPENMP +#include +#endif + +#ifdef __MPI +#include +#endif +namespace BSE +{ + +struct BlockHead +{ + int gr0; // global row start + int gc0; // global col start + int nr; // rows in block + int nc; // cols in block +}; + +#ifdef __MPI +inline MPI_Datatype mpi_type_blockhead() +{ + static MPI_Datatype dt = MPI_DATATYPE_NULL; + static bool committed = false; + if (!committed) + { + int blen[4] = {1, 1, 1, 1}; + MPI_Aint disp[4]; + MPI_Datatype types[4] = {MPI_INT, MPI_INT, MPI_INT, MPI_INT}; + + disp[0] = static_cast(offsetof(BlockHead, gr0)); + disp[1] = static_cast(offsetof(BlockHead, gc0)); + disp[2] = static_cast(offsetof(BlockHead, nr)); + disp[3] = static_cast(offsetof(BlockHead, nc)); + + MPI_Type_create_struct(4, blen, disp, types, &dt); + MPI_Type_commit(&dt); + committed = true; + } + return dt; +} +#endif + +/// @brief W[k_ai][k_bj] to 2d local matrix WA[aik1, bjk2] +template +void MolecularLRI::transform_k_2dlocal(std::vector& m_2d, + const std::map>>& m_lri, + const Parallel_2D& pm_2d, const double factor) +{ + ModuleBase::TITLE("MolecularLRI", "transform_k_2dlocal"); + ModuleBase::timer::start("MolecularLRI", "transform_k_2dlocal"); + const int npair = this->nocc * this->nvirt; + const double fac_base = factor * 2.0; // factor 2 for Ha → Ry; k-point weight multiplies below + const int nb = pm_2d.get_block_size(); +#ifdef __MPI + MPI_Datatype mpitype_blockhead = mpi_type_blockhead(); + // 0. outer loop of k1: communicate per 64 k1 + constexpr int comm_nk1 = 64; + std::vector send_head_counts(GlobalV::NPROC, 0), recv_head_counts(GlobalV::NPROC, 0); + std::vector send_buffer_counts(GlobalV::NPROC, 0), recv_buffer_counts(GlobalV::NPROC, 0); + std::vector send_buffer_counts_c(GlobalV::NPROC, 0); + std::vector shdispls(GlobalV::NPROC, 0), rhdispls(GlobalV::NPROC, 0); //displacements of block heads + std::vector sbdispls(GlobalV::NPROC, 0), rbdispls(GlobalV::NPROC, 0); //displacements of buffer + std::vector cursor_head(GlobalV::NPROC, 0), cursor_buffer(GlobalV::NPROC, 0); + std::vector send_heads, recv_heads; + std::vector send_buffers, recv_buffers; + int max_send_head_total = 0, max_recv_head_total = 0; + int max_send_buffer_total = 0, max_recv_buffer_total = 0; + for (int k1_start = 0; k1_start < this->nk; k1_start += comm_nk1) + { + std::fill(send_head_counts.begin(), send_head_counts.end(), 0); + std::fill(send_buffer_counts.begin(), send_buffer_counts.end(), 0); + std::fill(send_buffer_counts_c.begin(), send_buffer_counts_c.end(), 0); + std::fill(recv_head_counts.begin(), recv_head_counts.end(), 0); + std::fill(recv_buffer_counts.begin(), recv_buffer_counts.end(), 0); + const int k1_end = std::min(k1_start+comm_nk1, this->nk); + // 1. calculate and coummunicate block counts, then calculate block displacements + for (int kai = k1_start; kai < k1_end; ++kai) + { + if ( !this->is_local_k1[kai] ) continue; + const int row_base = kai * npair; + for (const int kbj : this->LR_lri.k2_indices) + { + const int col_base = kbj * npair; + for (int j = 0; j < npair; ) + { + const int global_col = col_base + j; + const int j_next = std::min(global_col/nb*nb + nb - col_base, npair); + for (int i = 0; i < npair; ) + { + const int global_row = row_base + i; + const int i_next = std::min(global_row/nb*nb + nb - row_base, npair); + if (!pm_2d.in_this_processor(global_row, global_col)) + { + const int owner = pm_2d.owner_processor(global_row, global_col); + ++send_head_counts[owner]; + send_buffer_counts_c[owner] += std::size_t(i_next - i) * std::size_t(j_next - j); + } + i = i_next; + } + j = j_next; + } + } + } + for (int p = 0; p < GlobalV::NPROC; ++p) + { + if (send_buffer_counts_c[p] > std::numeric_limits::max()) + { + throw std::overflow_error("in transform_k_2dlocal: overflow converting to int!"); + } + send_buffer_counts[p] = static_cast(send_buffer_counts_c[p]); + } + assert(send_head_counts.at(GlobalV::MY_RANK) == 0); + assert(send_buffer_counts.at(GlobalV::MY_RANK) == 0); + MPI_Alltoall(send_head_counts.data(), 1, MPI_INT, recv_head_counts.data(), 1, MPI_INT, pm_2d.comm()); + MPI_Alltoall(send_buffer_counts.data(), 1, MPI_INT, + recv_buffer_counts.data(), 1, MPI_INT, pm_2d.comm()); + + int send_head_total = 0, recv_head_total = 0, send_buffer_total = 0, recv_buffer_total = 0; + for (int p = 0; p < GlobalV::NPROC; ++p) + { + shdispls[p] = send_head_total; send_head_total += send_head_counts[p]; + rhdispls[p] = recv_head_total; recv_head_total += recv_head_counts[p]; + sbdispls[p] = send_buffer_total; send_buffer_total += send_buffer_counts[p]; + rbdispls[p] = recv_buffer_total; recv_buffer_total += recv_buffer_counts[p]; + } + if (send_head_total > max_send_head_total) + { max_send_head_total = send_head_total; send_heads.reserve(max_send_head_total); } + if (recv_head_total > max_recv_head_total) + { max_recv_head_total = recv_head_total; recv_heads.reserve(max_recv_head_total); } + if (send_buffer_total > max_send_buffer_total) + { max_send_buffer_total = send_buffer_total; send_buffers.reserve(max_send_buffer_total); } + if (recv_buffer_total > max_recv_buffer_total) + { max_recv_buffer_total = recv_buffer_total; recv_buffers.reserve(max_recv_buffer_total); } + + // 2. prepare block heads and buffers for send and recv + send_heads.resize(send_head_total); + recv_heads.resize(recv_head_total); + send_buffers.resize(send_buffer_total); + recv_buffers.resize(recv_buffer_total); + cursor_head = shdispls; + cursor_buffer = sbdispls; + for (int kai = k1_start; kai < k1_end; ++kai) + { + if (!this->is_local_k1[kai] ) continue; + const int row_base = kai * npair; + const double fac = fac_base * this->kv.wk[kai]; // k-point weight, normalized as sum = 1 + for (const int kbj : this->LR_lri.k2_indices) + { + const int col_base = kbj * npair; + const RI::Tensor& m_kai_kbj = m_lri.at(kai).at(kbj); + for (int j = 0; j < npair; ) + { + const int global_col = col_base + j; + const int j_next = std::min(global_col/nb*nb + nb - col_base, npair); + for (int i = 0; i < npair; ) + { + const int global_row = row_base + i; + const int i_next = std::min(global_row/nb*nb + nb - row_base, npair); + const int owner = pm_2d.owner_processor(global_row, global_col); + if (owner == GlobalV::MY_RANK) + { + const int lr0 = pm_2d.global2local_row(global_row); + const int lc0 = pm_2d.global2local_col(global_col); + const int lld = pm_2d.get_row_size(); + for(int jj = j; jj < j_next; ++jj) + { + for(int ii = i; ii < i_next; ++ii) + { + const int lr = lr0 + (ii - i); + const int lc = lc0 + (jj - j); + const std::size_t idx_2d = lr + std::size_t(lc) * std::size_t(lld); + m_2d[idx_2d] += (*m_kai_kbj.data)[ii + jj * npair] * fac; + } + } + } + else + { + BlockHead& head = send_heads[cursor_head[owner]++]; + head.gr0 = global_row; + head.gc0 = global_col; + head.nr = i_next - i; + head.nc = j_next - j; + for(int jj = j; jj < j_next; ++jj) + { + for (int ii = i; ii < i_next; ++ii) + { + send_buffers[cursor_buffer[owner]++] = (*m_kai_kbj.data)[ii + jj * npair] * fac; + } + } + } + i = i_next; + } + j = j_next; + } + } + } + for (int p = 0; p < GlobalV::NPROC; ++p) + { + assert(cursor_head[p] == shdispls[p] + send_head_counts[p]); + assert(cursor_buffer[p] == sbdispls[p] + send_buffer_counts[p]); + } + // 3. communicate block heads and buffers + MPI_Alltoallv(send_heads.data(), send_head_counts.data(), shdispls.data(), mpitype_blockhead, + recv_heads.data(), recv_head_counts.data(), rhdispls.data(), mpitype_blockhead, + pm_2d.comm()); + MPI_Alltoallv(send_buffers.data(), send_buffer_counts.data(), sbdispls.data(), LR_Util::MPIType::value(), + recv_buffers.data(), recv_buffer_counts.data(), rbdispls.data(), LR_Util::MPIType::value(), + pm_2d.comm()); + // 4. unpack recv head and buffer + int buf_cursor = 0; + for (int iblock = 0; iblock < recv_head_total; ++iblock) + { + const BlockHead& rh = recv_heads[iblock]; + const int nr = rh.nr; + const int nc = rh.nc; + const int lr = pm_2d.global2local_row(rh.gr0); + const int lc = pm_2d.global2local_col(rh.gc0); + const int lld = pm_2d.get_row_size(); + for (int j = 0; j < nc; ++j) + { + for (int i = 0; i < nr; ++i) + { + const int idx_buffer = i + j * nr; + const std::size_t idx_2d = (lr + i) + std::size_t(lc + j) * std::size_t(lld); + m_2d[idx_2d] += recv_buffers[idx_buffer + buf_cursor]; + } + } + buf_cursor += nr * nc; + } + assert(buf_cursor == recv_buffer_total); + } +#else + // gather all Wk + auto gather_matrix = [&](std::vector& target, + const std::valarray& value, + int k1_step, + int k2_step, + double factor) -> void + { + for (int j = 0; j < npair; ++j) + { + for (int i = 0; i < npair; ++i) + { + const std::size_t idx_target = (k1_step + i) + std::size_t(k2_step + j) * std::size_t(this->ndim); + const int idx_value = i + j * npair; + target[idx_target] += value[idx_value] * factor; + } + } + }; + + #ifdef _OPENMP + #pragma omp parallel for schedule(static) collapse(2) + #endif + for (const int kai : this->k1_indices) + { + const int k1_step = kai * npair; + for (const int kbj : this->k2_indices) + { + const int k2_step = kbj * npair; + const RI::Tensor& m_kai_kbj = m_lri.at(kai).at(kbj); + gather_matrix(m_2d, *m_kai_kbj.data, k1_step, k2_step, fac_base * this->kv.wk[kai]); + } + } +#endif + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "transform_k_2dlocal"); + ModuleBase::timer::end("MolecularLRI", "transform_k_2dlocal"); +} +} \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/test/CMakeLists.txt b/source/source_lcao/module_lr/bse/test/CMakeLists.txt new file mode 100644 index 0000000000..9126d8fd1f --- /dev/null +++ b/source/source_lcao/module_lr/bse/test/CMakeLists.txt @@ -0,0 +1,6 @@ +remove_definitions(-DUSE_LIBXC) +AddTest( + TARGET MODULE_LR_bse_solver_test + LIBS parameter base ${math_libs} ELPA::ELPA device psi + SOURCES hamilt_bse_solver_test.cpp ../hamilt_bse_solver.cpp +) \ No newline at end of file diff --git a/source/source_lcao/module_lr/bse/test/hamilt_bse_solver_test.cpp b/source/source_lcao/module_lr/bse/test/hamilt_bse_solver_test.cpp new file mode 100644 index 0000000000..bbd4137a34 --- /dev/null +++ b/source/source_lcao/module_lr/bse/test/hamilt_bse_solver_test.cpp @@ -0,0 +1,249 @@ +#include +#include "../hamilt_bse_solver.h" +#include +#include "source_base/module_container/base/third_party/blas.h" + +#define rand01 (static_cast(rand()) / static_cast(RAND_MAX) - 0.5 ) // [-0.5, 0.5] + +std::vector> generate_conjugate_matrix(int n) { + std::vector> matrix(n*n, 0.0); + for (int i = 0; i < n; ++i) { + for (int j = 0; j <= i; ++j) { + if(i==j) matrix[i * n + j] = std::complex(20+rand01, 0); // gaurantee {{A,B},{A*,B*}} is positive definite + else{ + matrix[i * n + j] = std::complex(rand01, rand01); + matrix[j * n + i] = std::conj(matrix[i * n + j]); + } + } + } + return matrix; +} +std::vector> generate_symmetry_matrix(int n) { + std::vector> matrix(n*n, 0.0); + for (int i = 0; i < n; ++i) { + for (int j = 0; j <= i; ++j) { + if(i==j) matrix[i * n + j] = std::complex(rand01, rand01); + else{ + matrix[i * n + j] = std::complex(rand01, rand01); + matrix[j * n + i] = matrix[i * n + j]; + } + } + } + return matrix; +} +void check_eq(std::complex* data1, std::complex* data2, int size, double eps) + { + for (int i = 0;i < size;++i) + { + EXPECT_NEAR(data1[i].real(), data2[i].real(), eps); + EXPECT_NEAR(data1[i].imag(), data2[i].imag(), eps); + } + }; + +TEST(BSETest, skewSolver) { + int my_rank, num_procs; + Cblacs_pinfo(&my_rank, &num_procs); + + int nA = 2; + std::vector ev(2*nA); + Parallel_2D pM, pA, pA_glb; + pM.init(2*nA, 2*nA, 1, MPI_COMM_WORLD, false); + pA.set(nA, nA, 1, pM.blacs_ctxt); + pA_glb.set(nA, nA, nA, pM.blacs_ctxt); + std::vector> v(pM.get_local_size(), 0.0); + std::vector> A_part(pA.get_local_size(), 0.0); + std::vector> B_part(pA.get_local_size(), 0.0); + std::vector> A_glb = { + {3.0, 0.0}, {0.5, -1.0}, {0.5, 1.0}, {6.0, 0.0} + }; + std::vector> B_glb = { + {1.2, 0.6}, {0.4, 0.5}, {0.4, 0.5}, {1.4, 0.3} + }; + Cpxgemr2d(nA, nA, A_glb.data(), 1, 1, pA_glb.desc, + A_part.data(), 1, 1, pA.desc, + pA_glb.blacs_ctxt); + Cpxgemr2d(nA, nA, B_glb.data(), 1, 1, pA_glb.desc, + B_part.data(), 1, 1, pA.desc, + pA_glb.blacs_ctxt); + BSE::solve_full(my_rank, A_part, B_part, pA, pM, ev, v); + EXPECT_NEAR(ev[0], -6.127295611, 1e-8); + EXPECT_NEAR(ev[1], -2.299184312, 1e-8); +} + +TEST(BSETest, TDASolver) { + int my_rank, num_procs; + Cblacs_pinfo(&my_rank, &num_procs); + int nA = 5; + Parallel_2D pA_glb, pA; + pA.init(nA, nA, 1, MPI_COMM_WORLD, false); + pA_glb.set(nA, nA, nA, pA.blacs_ctxt); + + std::vector ev(nA); + std::vector> A_glb; + if (my_rank == 0){ + A_glb = generate_conjugate_matrix(nA); + } + std::vector> A_part(pA.get_local_size(), 0.0); + std::vector> v(pA.get_local_size(), 0.0); + + Cpxgemr2d(nA, nA, A_glb.data(), 1, 1, pA_glb.desc, + A_part.data(), 1, 1, pA.desc, + pA_glb.blacs_ctxt); + + std::vector> Hv(pA.get_local_size(), 0.0); + std::vector> Ωv(pA.get_local_size(), 0.0); + + auto start = std::chrono::high_resolution_clock::now(); + BSE::solve_tda(my_rank, A_part, pA, ev, v); + auto end = std::chrono::high_resolution_clock::now(); + auto elapsed = std::chrono::duration_cast(end - start); + if (my_rank == 0){ + std::cout << "BSE::solve_tda execution time: " << elapsed.count() << " seconds" << std::endl; + } + + start = std::chrono::high_resolution_clock::now(); + #pragma omp parallel for collapse(2) + for (int i = 0; i < pA.get_col_size(); ++i) { + int col_glb = pA.local2global_col(i); + for (int j = 0; j < pA.get_row_size(); ++j) { + Ωv[i * pA.get_row_size() + j] = ev[col_glb] * v[i * pA.get_row_size() + j]; + } + } + end = std::chrono::high_resolution_clock::now(); + elapsed = std::chrono::duration_cast(end - start); + if (my_rank == 0) + std::cout << "Ωv construction time: " << elapsed.count() << " seconds" << std::endl; + + start = std::chrono::high_resolution_clock::now(); + ScalapackConnector::gemm('N', 'N', nA, nA, nA, 1.0, + A_part.data(), 1, 1, pA.desc, + v.data(), 1, 1, pA.desc, + 0.0, + Hv.data(), 1, 1, pA.desc); + end = std::chrono::high_resolution_clock::now(); + elapsed = std::chrono::duration_cast(end - start); + if (my_rank == 0) + std::cout << "Hv execution time: " << elapsed.count() << " seconds" << std::endl; + check_eq(Hv.data(), Ωv.data(), pA.get_local_size(), 1e-6); +} + +TEST(BSETest, skewSolver2) { + int my_rank, num_procs; + Cblacs_pinfo(&my_rank, &num_procs); + + int nA = 5; + int nM = 2 * nA; // Full matrix dimension + Parallel_2D pM, pA, pA_glb; + pM.init(nM, nM, 1, MPI_COMM_WORLD, false); + pA.set(nA, nA, 1, pM.blacs_ctxt); + pA_glb.set(nA, nA, nA, pM.blacs_ctxt); + + std::vector> A_glb, B_glb; + if (my_rank==0){ + A_glb = generate_conjugate_matrix(nA); + B_glb = generate_symmetry_matrix(nA); + } + std::vector> A_part(pA.get_local_size(), 0.0); + std::vector> B_part(pA.get_local_size(), 0.0); + std::vector ev(nM, 0.0); + std::vector> v(pM.get_local_size(), 0.0); + Cpxgemr2d(nA, nA, A_glb.data(), 1, 1, pA_glb.desc, + A_part.data(), 1, 1, pA.desc, + pA_glb.blacs_ctxt); + Cpxgemr2d(nA, nA, B_glb.data(), 1, 1, pA_glb.desc, + B_part.data(), 1, 1, pA.desc, + pA_glb.blacs_ctxt); + // solve_full destroys A_part and B_part; pass copies to preserve originals + auto A_copy = A_part; + auto B_copy = B_part; + auto start = std::chrono::high_resolution_clock::now(); + BSE::solve_full(my_rank, A_copy, B_copy, pA, pM, ev, v); + auto end = std::chrono::high_resolution_clock::now(); + auto elapsed = std::chrono::duration_cast(end - start); + if (my_rank == 0){ + std::cout << "BSE::solve_full execution time: " << elapsed.count() << " seconds" << std::endl; + } + + std::vector> full_H (pM.get_local_size(), 0.0); + std::vector> Hv(pM.get_local_size(), 0.0); + std::vector> Ωv(pM.get_local_size(), 0.0); + + start = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < pM.get_col_size(); ++i) { + int col_glb = pM.local2global_col(i); + for (int j = 0; j < pM.get_row_size(); ++j) { + Ωv[i * pM.get_row_size() + j] = ev[col_glb] * v[i * pM.get_row_size() + j]; + } + } + end = std::chrono::high_resolution_clock::now(); + elapsed = std::chrono::duration_cast(end - start); + if (my_rank == 0) + std::cout << "Ωv construction time: " << elapsed.count() << " seconds" << std::endl; + + BSE::arrayFlatten1(A_part, B_part, full_H, pA, pM); + + start = std::chrono::high_resolution_clock::now(); + ScalapackConnector::gemm('N', 'N', nM, nM, nM, 1.0, + full_H.data(), 1, 1, pM.desc, + v.data(), 1, 1, pM.desc, + 0.0, + Hv.data(), 1, 1, pM.desc); + end = std::chrono::high_resolution_clock::now(); + elapsed = std::chrono::duration_cast(end - start); + if (my_rank == 0) + std::cout << "Hv execution time: " << elapsed.count() << " seconds" << std::endl; + + check_eq(Hv.data(), Ωv.data(), pM.get_local_size(), 1e-6); + + std::vector> identity(pM.get_local_size(), 0.0); + for (int i = 0; i < pM.get_col_size(); ++i) { + int col_glb = pM.local2global_col(i); + int row_loc = pM.global2local_row(col_glb); + if (row_loc != -1) + identity[i * pM.get_row_size() + row_loc] = 1.0; + } + // v = ｢Y* X left_v = ｢-Y* X + // X* Y｣ X* -Y｣ + std::vector> left_v(pM.get_local_size(), 0.0); + for (int i = 0; i < pM.get_col_size(); ++i) { + int col_glb = pM.local2global_col(i); + for (int j = 0; j < pM.get_row_size(); ++j) { + int row_glb = pM.local2global_row(j); + if (col_glb < nA && row_glb < nA) + left_v[i * pM.get_row_size() + j] = -v[i * pM.get_row_size() + j]; + else if (col_glb >= nA && row_glb < nA) + left_v[i * pM.get_row_size() + j] = v[i * pM.get_row_size() + j]; + else if (col_glb < nA && row_glb >= nA) + left_v[i * pM.get_row_size() + j] = v[i * pM.get_row_size() + j]; + else // if (col_glb >= nA && row_glb >= nA) + left_v[i * pM.get_row_size() + j] = -v[i * pM.get_row_size() + j]; + } + } + + start = std::chrono::high_resolution_clock::now(); + ScalapackConnector::gemm('C', 'N', nM, nM, nM, 1.0, + left_v.data(), 1, 1, pM.desc, + v.data(), 1, 1, pM.desc, + 0.0, + Ωv.data(), 1, 1, pM.desc);// overwrite Ωv by left_v.v + end = std::chrono::high_resolution_clock::now(); + elapsed = std::chrono::duration_cast(end - start); + if (my_rank == 0) + std::cout << "left_v.v execution time: " << elapsed.count() << " seconds" << std::endl; + check_eq(Ωv.data(), identity.data(), pM.get_local_size(), 1e-6); +} + + +int main(int argc, char **argv) { + srand(time(nullptr)); + int thread_level = -1; + MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &thread_level); + if (thread_level != MPI_THREAD_MULTIPLE) + { + std::cerr << "MPI_Init_thread request " << MPI_THREAD_MULTIPLE << " but provide " << thread_level << std::endl; + } + testing::InitGoogleTest(&argc, argv); + int result = RUN_ALL_TESTS(); + MPI_Finalize(); + return result; +} \ No newline at end of file diff --git a/source/source_lcao/module_lr/dm_trans/dm_trans.h b/source/source_lcao/module_lr/dm_trans/dm_trans.h index dfdf11cabf..36bbd30877 100644 --- a/source/source_lcao/module_lr/dm_trans/dm_trans.h +++ b/source/source_lcao/module_lr/dm_trans/dm_trans.h @@ -2,6 +2,7 @@ // use tensor or basematrix in the future #include #include "source_psi/psi.h" +#include "source_lcao/module_lr/utils/mo_type.h" #include #ifdef __MPI #include "source_base/parallel_2d.h" @@ -9,11 +10,6 @@ namespace LR { -#ifndef MO_TYPE_H -#define MO_TYPE_H - enum MO_TYPE { OO, VO, VV }; -#endif - #ifdef __MPI /// @brief calculate the 2d-block transition density matrix in AO basis using p?gemm /// \f[ \tilde{\rho}_{\mu_j\mu_b}=\sum_{jb}c_{j,\mu_j}X_{jb}c^*_{b,\mu_b} \f] @@ -28,7 +24,7 @@ namespace LR const int nvirt, const Parallel_2D& pmat, const T factor = (T)1.0, - const MO_TYPE type = MO_TYPE::VO); + const LR_Util::MO_TYPE type = LR_Util::MO_TYPE::VO); #endif /// @brief calculate the 2d-block transition density matrix in AO basis using ?gemm @@ -38,7 +34,7 @@ namespace LR const psi::Psi& c, const int& nocc, const int& nvirt, const T factor = (T)1.0, - const MO_TYPE type = MO_TYPE::VO); + const LR_Util::MO_TYPE type = LR_Util::MO_TYPE::VO); // for test /// @brief calculate the 2d-block transition density matrix in AO basis using for loop (for test) @@ -48,5 +44,5 @@ namespace LR const psi::Psi& c, const int& nocc, const int& nvirt, const T factor = (T)1.0, - const MO_TYPE type = MO_TYPE::VO); + const LR_Util::MO_TYPE type = LR_Util::MO_TYPE::VO); } \ No newline at end of file diff --git a/source/source_lcao/module_lr/dm_trans/dm_trans_parallel.cpp b/source/source_lcao/module_lr/dm_trans/dm_trans_parallel.cpp index 5b36f798b2..f8198afcaf 100644 --- a/source/source_lcao/module_lr/dm_trans/dm_trans_parallel.cpp +++ b/source/source_lcao/module_lr/dm_trans/dm_trans_parallel.cpp @@ -19,7 +19,7 @@ std::vector cal_dm_trans_pblas(const double* const X_istate, const int nvirt, const Parallel_2D& pmat, const double factor, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_pblas"); assert(px.comm() == pc.comm() && px.comm() == pmat.comm()); @@ -28,11 +28,12 @@ std::vector cal_dm_trans_pblas(const double* const X_istate, const int nks = c.get_nk(); const int i1 = 1; - const int ivirt = nocc + 1; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? ivirt : i1; - const int imo2 = type == MO_TYPE::OO ? i1 : ivirt; + int nmo1_set, nmo2_set, imo1_set, imo2_set; + LR_Util::set_dim(type, nocc, nvirt, nmo1_set, nmo2_set, imo1_set, imo2_set); + const int nmo1 = nmo1_set; + const int nmo2 = nmo2_set; + const int imo1 = imo1_set + 1; + const int imo2 = imo2_set + 1; std::vector dm_trans(nks, container::Tensor(DAT::DT_DOUBLE, DEV::CpuDevice, { pmat.get_col_size(), pmat.get_row_size() })); @@ -76,7 +77,7 @@ std::vector cal_dm_trans_pblas(const std::complex* co const int nvirt, const Parallel_2D& pmat, const std::complex factor, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_pblas"); assert(px.comm() == pc.comm() && px.comm() == pmat.comm()); @@ -84,11 +85,12 @@ std::vector cal_dm_trans_pblas(const std::complex* co assert(pmat.get_local_size() > 0); const int nks = c.get_nk(); const int i1 = 1; - const int ivirt = nocc + 1; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; - const int imo1 = type == MO_TYPE::VV ? ivirt : i1; - const int imo2 = type == MO_TYPE::OO ? i1 : ivirt; + int nmo1_set, nmo2_set, imo1_set, imo2_set; + LR_Util::set_dim(type, nocc, nvirt, nmo1_set, nmo2_set, imo1_set, imo2_set); + const int nmo1 = nmo1_set; + const int nmo2 = nmo2_set; + const int imo1 = imo1_set + 1; + const int imo2 = imo2_set + 1; std::vector dm_trans(nks, container::Tensor(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, {pmat.get_col_size(), pmat.get_row_size()})); diff --git a/source/source_lcao/module_lr/dm_trans/dm_trans_serial.cpp b/source/source_lcao/module_lr/dm_trans/dm_trans_serial.cpp index 3a4553e0b4..521e492189 100644 --- a/source/source_lcao/module_lr/dm_trans/dm_trans_serial.cpp +++ b/source/source_lcao/module_lr/dm_trans/dm_trans_serial.cpp @@ -11,15 +11,13 @@ namespace LR const int& nocc, const int& nvirt, const double factor, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { // cxc_out_test(X_istate, c); ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_forloop"); const int nks = c.get_nk(); - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); const int naos = c.get_nbasis(); std::vector dm_trans(nks, container::Tensor(DAT::DT_DOUBLE, DEV::CpuDevice, { naos, naos })); @@ -51,16 +49,14 @@ namespace LR const int& nocc, const int& nvirt, const std::complex factor, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { // cxc_out_test(X_istate, c); ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_forloop"); const int nks = c.get_nk(); const int naos = c.get_nbasis(); - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); std::vector dm_trans(nks, container::Tensor(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, { naos, naos })); for (auto& dm : dm_trans)ModuleBase::GlobalFunc::ZEROS(dm.data>(), naos * naos); // loop for AOs @@ -92,15 +88,13 @@ namespace LR const int& nocc, const int& nvirt, const double factor, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_blas"); const int nks = c.get_nk(); const int naos = c.get_nbasis(); - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); std::vector dm_trans(nks, container::Tensor(DAT::DT_DOUBLE, DEV::CpuDevice, { naos, naos })); for (size_t isk = 0;isk < nks;++isk) { @@ -130,15 +124,13 @@ namespace LR const int& nocc, const int& nvirt, const std::complex factor, - const MO_TYPE type) + const LR_Util::MO_TYPE type) { ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_blas"); const int nks = c.get_nk(); const int naos = c.get_nbasis(); - const int imo1 = type == MO_TYPE::VV ? nocc : 0; - const int imo2 = type == MO_TYPE::OO ? 0 : nocc; - const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc; - const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt; + int nmo1, nmo2, imo1, imo2; + LR_Util::set_dim(type, nocc, nvirt, nmo1, nmo2, imo1, imo2); std::vector dm_trans(nks, container::Tensor(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, { naos, naos })); for (size_t isk = 0;isk < nks;++isk) { diff --git a/source/source_lcao/module_lr/dm_trans/test/dm_trans_test.cpp b/source/source_lcao/module_lr/dm_trans/test/dm_trans_test.cpp index ce4bde3c6b..53b16661b6 100644 --- a/source/source_lcao/module_lr/dm_trans/test/dm_trans_test.cpp +++ b/source/source_lcao/module_lr/dm_trans/test/dm_trans_test.cpp @@ -74,16 +74,16 @@ TEST_F(DMTransTest, DoubleSerial) std::vector temp(s.nks, s.naos); psi::Psi c(s.nks, s.nocc + s.nvirt, s.naos, temp, true); set_rand(c.get_pointer(), size_c); - auto test = [&](psi::Psi& X, const LR::MO_TYPE type) + auto test = [&](psi::Psi& X, const LR_Util::MO_TYPE type) { X.fix_b(istate); const std::vector& dm_for = LR::cal_dm_trans_forloop_serial(X.get_pointer(), c, s.nocc, s.nvirt, 1., type); const std::vector& dm_blas = LR::cal_dm_trans_blas(X.get_pointer(), c, s.nocc, s.nvirt, 1., type); for (int isk = 0;isk < s.nks;++isk) check_eq(dm_for[isk].data(), dm_blas[isk].data(), s.naos * s.naos); }; - test(X_vo, LR::MO_TYPE::VO); - test(X_oo, LR::MO_TYPE::OO); - test(X_vv, LR::MO_TYPE::VV); + test(X_vo, LR_Util::MO_TYPE::VO); + test(X_oo, LR_Util::MO_TYPE::OO); + test(X_vv, LR_Util::MO_TYPE::VV); } } @@ -105,16 +105,16 @@ TEST_F(DMTransTest, ComplexSerial) std::vector temp(s.nks, s.naos); psi::Psi> c(s.nks, s.nocc + s.nvirt, s.naos, temp, true); set_rand(c.get_pointer(), size_c); - auto test = [&](psi::Psi>& X, const LR::MO_TYPE type) + auto test = [&](psi::Psi>& X, const LR_Util::MO_TYPE type) { X.fix_b(istate); const std::vector& dm_for = LR::cal_dm_trans_forloop_serial(X.get_pointer(), c, s.nocc, s.nvirt, std::complex(1., 0.), type); const std::vector& dm_blas = LR::cal_dm_trans_blas(X.get_pointer(), c, s.nocc, s.nvirt, std::complex(1., 0.), type); for (int isk = 0;isk < s.nks;++isk) check_eq(dm_for[isk].data>(), dm_blas[isk].data>(), s.naos * s.naos); }; - test(X_vo, LR::MO_TYPE::VO); - test(X_oo, LR::MO_TYPE::OO); - test(X_vv, LR::MO_TYPE::VV); + test(X_vo, LR_Util::MO_TYPE::VO); + test(X_oo, LR_Util::MO_TYPE::OO); + test(X_vv, LR_Util::MO_TYPE::VV); } } @@ -159,6 +159,7 @@ TEST_F(DMTransTest, DoubleParallel) auto gather = [&](const psi::Psi& X, psi::Psi& X_full, const Parallel_2D& px, const int dim1, const int dim2) { + X_full.zero_out(); for (int istate = 0;istate < nstate;++istate) { X.fix_b(istate); @@ -183,6 +184,7 @@ TEST_F(DMTransTest, DoubleParallel) // gather C std::vector temp(s.nks, s.naos); psi::Psi c_full(s.nks, s.nocc + s.nvirt, s.naos, temp, true); + c_full.zero_out(); for (int isk = 0;isk < s.nks;++isk) { c.fix_k(isk); @@ -190,7 +192,7 @@ TEST_F(DMTransTest, DoubleParallel) LR_Util::gather_2d_to_full(pc, c.get_pointer(), c_full.get_pointer(), false, s.naos, s.nocc + s.nvirt); } - auto test = [&](psi::Psi& X, psi::Psi& X_full, const Parallel_2D& px, const LR::MO_TYPE type) + auto test = [&](psi::Psi& X, psi::Psi& X_full, const Parallel_2D& px, const LR_Util::MO_TYPE type) { X.fix_b(istate); X_full.fix_b(istate); @@ -198,6 +200,7 @@ TEST_F(DMTransTest, DoubleParallel) std::vector dm_gather(s.nks, container::Tensor(DAT::DT_DOUBLE, DEV::CpuDevice, { s.naos, s.naos })); for (int isk = 0;isk < s.nks;++isk) { + dm_gather[isk].zero(); LR_Util::gather_2d_to_full(pmat, dm_pblas_loc[isk].data(), dm_gather[isk].data(), false, s.naos, s.naos); } if (my_rank == 0) @@ -206,9 +209,9 @@ TEST_F(DMTransTest, DoubleParallel) for (int isk = 0;isk < s.nks;++isk) check_eq(dm_full[isk].data(), dm_gather[isk].data(), s.naos * s.naos); } }; - test(X_vo, X_full_vo, px_vo, LR::MO_TYPE::VO); - test(X_oo, X_full_oo, px_oo, LR::MO_TYPE::OO); - test(X_vv, X_full_vv, px_vv, LR::MO_TYPE::VV); + test(X_vo, X_full_vo, px_vo, LR_Util::MO_TYPE::VO); + test(X_oo, X_full_oo, px_oo, LR_Util::MO_TYPE::OO); + test(X_vv, X_full_vv, px_vv, LR_Util::MO_TYPE::VV); } } } @@ -244,6 +247,7 @@ TEST_F(DMTransTest, ComplexParallel) auto gather = [&](const psi::Psi>& X, psi::Psi>& X_full, const Parallel_2D& px, const int dim1, const int dim2) { + X_full.zero_out(); for (int istate = 0;istate < nstate;++istate) { X.fix_b(istate); @@ -267,6 +271,7 @@ TEST_F(DMTransTest, ComplexParallel) // compare to global matrix std::vector ngk_temp_2(s.nks, s.naos); psi::Psi> c_full(s.nks, s.nocc + s.nvirt, s.naos, ngk_temp_2, true); + c_full.zero_out(); for (int isk = 0;isk < s.nks;++isk) { c.fix_k(isk); @@ -274,7 +279,7 @@ TEST_F(DMTransTest, ComplexParallel) LR_Util::gather_2d_to_full(pc, c.get_pointer(), c_full.get_pointer(), false, s.naos, s.nocc + s.nvirt); } - auto test = [&](psi::Psi>& X, psi::Psi>& X_full, const Parallel_2D& px, const LR::MO_TYPE type) + auto test = [&](psi::Psi>& X, psi::Psi>& X_full, const Parallel_2D& px, const LR_Util::MO_TYPE type) { X.fix_b(istate); X_full.fix_b(istate); @@ -282,6 +287,7 @@ TEST_F(DMTransTest, ComplexParallel) std::vector dm_gather(s.nks, container::Tensor(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, { s.naos, s.naos })); for (int isk = 0;isk < s.nks;++isk) { + dm_gather[isk].zero(); LR_Util::gather_2d_to_full(pmat, dm_pblas_loc[isk].data>(), dm_gather[isk].data>(), false, s.naos, s.naos); } if (my_rank == 0) @@ -290,9 +296,9 @@ TEST_F(DMTransTest, ComplexParallel) for (int isk = 0;isk < s.nks;++isk) check_eq(dm_full[isk].data>(), dm_gather[isk].data>(), s.naos * s.naos); } }; - test(X_vo, X_full_vo, px_vo, LR::MO_TYPE::VO); - test(X_oo, X_full_oo, px_oo, LR::MO_TYPE::OO); - test(X_vv, X_full_vv, px_vv, LR::MO_TYPE::VV); + test(X_vo, X_full_vo, px_vo, LR_Util::MO_TYPE::VO); + test(X_oo, X_full_oo, px_oo, LR_Util::MO_TYPE::OO); + test(X_vv, X_full_vv, px_vv, LR_Util::MO_TYPE::VV); } } } diff --git a/source/source_lcao/module_lr/hamilt_casida.cpp b/source/source_lcao/module_lr/hamilt_casida.cpp index 35238e178f..e51ddee793 100644 --- a/source/source_lcao/module_lr/hamilt_casida.cpp +++ b/source/source_lcao/module_lr/hamilt_casida.cpp @@ -10,6 +10,10 @@ namespace LR const int nv = this->nvirt[0]; const auto& px = this->pX[0]; const int ldim = nk * px.get_local_size(); + if (px.get_local_size() == 0) { + std::cerr<< "Warning: Parallel_2D in RANK "+std::to_string(GlobalV::MY_RANK) +" has no local size, please use less mpi." << std::endl; + std::cerr<< " [File:"<<__FILE__<< ", Function: " << __FUNCTION__ << ", Line: " << __LINE__ << "]" << std::endl; + } int npairs = no * nv; std::vector Amat_full(this->nk * npairs * this->nk * npairs, 0.0); for (int ik = 0;ik < this->nk;++ik) { @@ -45,7 +49,9 @@ namespace LR LR_Util::gather_2d_to_full(px, &A_aibj.get_pointer()[ik_ai * px.get_local_size()], Amat_full.data() + kbj * this->nk * npairs /*col, bj*/ + ik_ai * npairs/*row, ai*/, false, nv, no); -} + } +#else + std::memcpy(Amat_full.data() + kbj * this->nk * npairs, A_aibj.get_pointer(), this->nk * npairs * sizeof(T)); #endif } } diff --git a/source/source_lcao/module_lr/hamilt_casida.h b/source/source_lcao/module_lr/hamilt_casida.h index d835fad2d3..b9459cfa8a 100644 --- a/source/source_lcao/module_lr/hamilt_casida.h +++ b/source/source_lcao/module_lr/hamilt_casida.h @@ -10,6 +10,7 @@ #include "source_lcao/module_lr/operator_casida/operator_lr_exx.h" #include "source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h" #include "source_lcao/module_ri/LRI_CV_Tools.h" +#include "source_lcao/module_lr/ri_benchmark/ri_benchmark.h" #endif namespace LR { @@ -42,15 +43,15 @@ namespace LR : nspin(nspin), nocc(nocc), nvirt(nvirt), pX(pX_in), nk(kv_in.get_nks() / nspin) { ModuleBase::TITLE("HamiltLR", "HamiltLR"); - if (ri_hartree_benchmark != "aims") { assert(aims_nbasis.empty()); } + if (ri_hartree_benchmark != "aims" && ri_hartree_benchmark !="aims-librpa") { assert(aims_nbasis.empty()); } // always use nspin=1 for transition density matrix this->DM_trans = LR_Util::make_unique>(&pmat_in, 1, kv_in.kvec_d, nk); if (ri_hartree_benchmark == "none") { LR_Util::initialize_DMR(*this->DM_trans, pmat_in, ucell_in, gd_in, orb_cutoff); } // this->DM_trans->init_DMR(&gd_in, &ucell_in); // too large due to not restricted by orb_cutoff - // add the diag operator (the first one) + // 1.add the diag operator (the first one) this->ops = new OperatorLRDiag(eig_ks.c, pX[0], nk, nocc[0], nvirt[0]); - //add Hxc operator + // 2.add Hxc operator #ifdef __EXX using TAC = std::pair>; using TLRI = std::map>>; @@ -68,23 +69,33 @@ namespace LR #ifdef __EXX if (spin_type == "singlet") { - if (ri_hartree_benchmark == "aims") + if (ri_hartree_benchmark == "aims" || ri_hartree_benchmark == "aims-librpa") { - Cs_read = LRI_CV_Tools::read_Cs_ao(dir + "Cs_data_0.txt"); - Vs_read = RI_Benchmark::read_coulomb_mat_general(dir + "coulomb_mat_0.txt", Cs_read); + LR_IO::RI_kRlist kRlist (ucell_in, const_cast(&kv_in)); + // though C and V are real, here still use to multiply with psi + Cs_read = LRI_CV_Tools::read_Cs_ao_all(dir); + Vs_read = LR_IO::read_coulomb_mat_general_k(dir, Cs_read, kRlist); } else if (ri_hartree_benchmark == "abacus") { Cs_read = LRI_CV_Tools::read_Cs_ao(dir + "Cs"); Vs_read = LRI_CV_Tools::read_Vs_abf(dir + "Vs"); } - if (!std::set({ "rpa", "hf" }).count(xc_kernel)) { throw std::runtime_error("ri_hartree_benchmark is only supported for xc_kernel rpa and hf"); } + else if (ri_hartree_benchmark == "abacus-librpa") + { + LR_IO::RI_kRlist kRlist (ucell_in, const_cast(&kv_in)); + Cs_read = LRI_CV_Tools::read_Cs_ao_all(dir); + Vs_read = LR_IO::read_coulomb_mat_k(dir, Cs_read, kRlist); + } + if (!std::set({ "rpa", "hf"}).count(xc_kernel)) { + throw std::runtime_error("ri_hartree_benchmark is only supported for xc_kernel = rpa, hf"); + } RI_Benchmark::OperatorRIHartree* ri_hartree_op = new RI_Benchmark::OperatorRIHartree(ucell_in, naos, nocc[0], nvirt[0], psi_ks_in, - Cs_read, Vs_read, ri_hartree_benchmark == "aims", aims_nbasis); + Cs_read, Vs_read); this->ops->add(ri_hartree_op); } - else if (spin_type == "triplet") { std::cout << "f_Hxc based on grid integral is not needed." << std::endl; } + else if (spin_type == "triplet") { std::cout << "Hatree term is not needed for S2:triplet." << std::endl; } #else ModuleBase::WARNING_QUIT("ESolver_LR", "RI benchmark is only supported when compile with LibRI."); #endif @@ -96,9 +107,9 @@ namespace LR this->DM_trans, pot_in, ucell_in, orb_cutoff, gd_in, kv_in, pX_in, pc_in, pmat_in); this->ops->add(lr_hxc); } -#ifdef __EXX +#ifdef __EXX// 3.add Exx operator if (xc_kernel == "hf" || xc_kernel == "hse") - { //add Exx operator + { if (ri_hartree_benchmark != "none" && spin_type == "singlet") { exx_lri_in.lock()->reset_Cs(Cs_read); @@ -107,8 +118,7 @@ namespace LR // std::cout << "exx_alpha=" << exx_alpha << std::endl; // the default value of exx_alpha is 0.25 when dft_functional is pbe or hse hamilt::Operator* lr_exx = new OperatorLREXX(nspin, naos, nocc[0], nvirt[0], ucell_in, psi_ks_in, this->DM_trans, exx_lri_in, kv_in, pX_in[0], pc_in, pmat_in, - xc_kernel == "hf" ? 1.0 : exx_alpha, //alpha - aims_nbasis); + (xc_kernel == "hf") ? 1.0 : exx_alpha); this->ops->add(lr_exx); } #endif @@ -148,29 +158,29 @@ namespace LR } } - void global2local(T* lvec, const T* gvec, const int& nband) const - { - const int npairs = nocc[0] * nvirt[0]; - for (int ib = 0;ib < nband;++ib) - { - const int loffset_b = ib * nk * pX[0].get_local_size(); - const int goffset_b = ib * nk * npairs; - for (int ik = 0;ik < nk;++ik) - { - const int loffset = loffset_b + ik * pX[0].get_local_size(); - const int goffset = goffset_b + ik * npairs; - for (int lo = 0;lo < pX[0].get_col_size();++lo) - { - const int go = pX[0].local2global_col(lo); - for (int lv = 0;lv < pX[0].get_row_size();++lv) - { - const int gv = pX[0].local2global_row(lv); - lvec[loffset + lo * pX[0].get_row_size() + lv] = gvec[goffset + go * nvirt[0] + gv]; - } - } - } - } - } + // void global2local(T* lvec, const T* gvec, const int& nband) const + // { + // const int npairs = nocc[0] * nvirt[0]; + // for (int ib = 0;ib < nband;++ib) + // { + // const int loffset_b = ib * nk * pX[0].get_local_size(); + // const int goffset_b = ib * nk * npairs; + // for (int ik = 0;ik < nk;++ik) + // { + // const int loffset = loffset_b + ik * pX[0].get_local_size(); + // const int goffset = goffset_b + ik * npairs; + // for (int lo = 0;lo < pX[0].get_col_size();++lo) + // { + // const int go = pX[0].local2global_col(lo); + // for (int lv = 0;lv < pX[0].get_row_size();++lv) + // { + // const int gv = pX[0].local2global_row(lv); + // lvec[loffset + lo * pX[0].get_row_size() + lv] = gvec[goffset + go * nvirt[0] + gv]; + // } + // } + // } + // } + // } private: const std::vector& nocc; diff --git a/source/source_lcao/module_lr/hamilt_ulr.hpp b/source/source_lcao/module_lr/hamilt_ulr.hpp index 4f5fdfbfd9..38d77e753b 100644 --- a/source/source_lcao/module_lr/hamilt_ulr.hpp +++ b/source/source_lcao/module_lr/hamilt_ulr.hpp @@ -118,7 +118,7 @@ namespace LR const std::vector npairs = { this->nocc[0] * this->nvirt[0], this->nocc[1] * this->nvirt[1] }; const std::vector ldim_is = { nk * pX[0].get_local_size(), nk * pX[1].get_local_size() }; const std::vector gdim_is = { nk * npairs[0], nk * npairs[1] }; - std::vector Amat_full(gdim * gdim); + std::vector Amat_full(gdim * gdim, 0.0); for (int is_bj : {0, 1}) { const int no = this->nocc[is_bj]; @@ -172,36 +172,36 @@ namespace LR } /// copy global data (eigenvectors) to local memory - void global2local(T* lvec, const T* gvec, const int& nband) const - { - const std::vector npairs = { this->nocc[0] * this->nvirt[0], this->nocc[1] * this->nvirt[1] }; - const std::vector ldim_is = { nk * pX[0].get_local_size(), nk * pX[1].get_local_size() }; - const std::vector gdim_is = { nk * npairs[0], nk * npairs[1] }; - for (int ib = 0;ib < nband;++ib) - { - const int loffset_b = ib * this->ldim; - const int goffset_b = ib * this->gdim; - for (int is : {0, 1}) - { - const int loffset_bs = loffset_b + is * ldim_is[0]; - const int goffset_bs = goffset_b + is * gdim_is[0]; - for (int ik = 0;ik < nk;++ik) - { - const int loffset = loffset_bs + ik * pX[is].get_local_size(); - const int goffset = goffset_bs + ik * npairs[is]; - for (int lo = 0;lo < pX[is].get_col_size();++lo) - { - const int go = pX[is].local2global_col(lo); - for (int lv = 0;lv < pX[is].get_row_size();++lv) - { - const int gv = pX[is].local2global_row(lv); - lvec[loffset + lo * pX[is].get_row_size() + lv] = gvec[goffset + go * nvirt[is] + gv]; - } - } - } - } - } - } + // void global2local(T* lvec, const T* gvec, const int& nband) const + // { + // const std::vector npairs = { this->nocc[0] * this->nvirt[0], this->nocc[1] * this->nvirt[1] }; + // const std::vector ldim_is = { nk * pX[0].get_local_size(), nk * pX[1].get_local_size() }; + // const std::vector gdim_is = { nk * npairs[0], nk * npairs[1] }; + // for (int ib = 0;ib < nband;++ib) + // { + // const int loffset_b = ib * this->ldim; + // const int goffset_b = ib * this->gdim; + // for (int is : {0, 1}) + // { + // const int loffset_bs = loffset_b + is * ldim_is[0]; + // const int goffset_bs = goffset_b + is * gdim_is[0]; + // for (int ik = 0;ik < nk;++ik) + // { + // const int loffset = loffset_bs + ik * pX[is].get_local_size(); + // const int goffset = goffset_bs + ik * npairs[is]; + // for (int lo = 0;lo < pX[is].get_col_size();++lo) + // { + // const int go = pX[is].local2global_col(lo); + // for (int lv = 0;lv < pX[is].get_row_size();++lv) + // { + // const int gv = pX[is].local2global_row(lv); + // lvec[loffset + lo * pX[is].get_row_size() + lv] = gvec[goffset + go * nvirt[is] + gv]; + // } + // } + // } + // } + // } + // } private: const std::vector& nocc; diff --git a/source/source_lcao/module_lr/hsolver_lrtd.hpp b/source/source_lcao/module_lr/hsolver_lrtd.hpp index 4d82bc6a0a..24a36b59a5 100644 --- a/source/source_lcao/module_lr/hsolver_lrtd.hpp +++ b/source/source_lcao/module_lr/hsolver_lrtd.hpp @@ -4,7 +4,6 @@ #include "source_hsolver/diago_dav_subspace.h" #include "source_hsolver/diago_cg.h" #include "source_hsolver/diago_iter_assist.h" -#include "source_hsolver/diago_cg.h" #include "source_lcao/module_lr/utils/lr_util.h" #include "source_lcao/module_lr/utils/lr_util_print.h" #include "source_base/module_container/ATen/core/tensor_map.h" @@ -29,6 +28,10 @@ namespace LR T* psi, const int& dim, ///< local leading dimension (or nbasis) const int& nband, ///< nstates in LR-TDDFT, not (nocc+nvirt) + const int& nk, + const std::vector& nocc, + const std::vector& nvirt, + const std::vector& pX, double* eig, const std::string method, const Real& diag_ethr, ///< threshold for diagonalization @@ -62,8 +65,14 @@ namespace LR print_eigs(eig_complex, "Right eigenvalues: of the non-Hermitian matrix: (Ry)"); for (int i = 0; i < gdim; i++) { eigenvalue[i] = eig_complex[i].real(); } } + bool openshell = std::is_same>::value; // copy eigenvectors - hm.global2local(psi, Amat_full.data(), nband); +#ifdef __MPI + LR_Util::global2local_X(psi, Amat_full.data(), nband, nk, + nocc, nvirt, pX, openshell); +#else + std::memcpy(psi, Amat_full.data(), sizeof(T) * nband * gdim); +#endif } else { diff --git a/source/source_lcao/module_lr/lr_spectrum.cpp b/source/source_lcao/module_lr/lr_spectrum.cpp index e0aa8aa6b7..4ad84d8481 100644 --- a/source/source_lcao/module_lr/lr_spectrum.cpp +++ b/source/source_lcao/module_lr/lr_spectrum.cpp @@ -19,10 +19,10 @@ elecstate::DensityMatrix LR::LR_Spectrum::cal_transition_density_matrix const int offset_x = offset_b + is * nk * this->pX[0].get_local_size(); //1. transition density #ifdef __MPI - std::vector dm_trans_2d = cal_dm_trans_pblas(X + offset_x, this->pX[is], psi_ks[is], this->pc, this->naos, this->nocc[is], this->nvirt[is], this->pmat, (T)1.0 / (T)nk); + std::vector dm_trans_2d = cal_dm_trans_pblas(X + offset_x, this->pX[is], psi_ks_vec[is], this->pc, this->naos, this->nocc[is], this->nvirt[is], this->pmat, (T)1.0 / (T)nk); // if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos, pmat); #else - std::vector dm_trans_2d = cal_dm_trans_blas(X + offset_x, this->psi_ks[is], this->nocc[is], this->nvirt[is], (T)1.0 / (T)nk); + std::vector dm_trans_2d = cal_dm_trans_blas(X + offset_x, this->psi_ks_vec[is], this->nocc[is], this->nvirt[is], (T)1.0 / (T)nk); // if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos); #endif for (int ik = 0;ik < this->nk;++ik) { DM_trans.set_DMK_pointer(ik + is * nk, dm_trans_2d[ik].data()); } @@ -37,11 +37,12 @@ elecstate::DensityMatrix LR::LR_Spectrum::cal_transition_density_matrix inline void check_sum_rule(const double& osc_tot) { + GlobalV::ofs_running << "Total oscillator strength = " << osc_tot << std::endl; + std::cout << "Total oscillator strength = " << osc_tot << std::endl; if (std::abs(osc_tot - 1.0) > 1e-3) { - GlobalV::ofs_running << "Warning: in LR_Spectrum::oscillator_strength, \ - the sum rule is not satisfied, try more nstates if needed.\n \ - Total oscillator strength = " + std::to_string(osc_tot) + "\n"; -} + GlobalV::ofs_running << "The sum rule is not well satisfied, try more nvirt and nstates if needed." << std::endl; + std::cout << "The sum rule is not well satisfied, try more nvirt and nstates if needed." << std::endl; + } } template<> @@ -49,7 +50,7 @@ ModuleBase::Vector3 LR::LR_Spectrum::cal_transition_dipole_istat { ModuleBase::Vector3 trans_dipole(0.0, 0.0, 0.0); // 1. transition density matrix - const elecstate::DensityMatrix& DM_trans = this->cal_transition_density_matrix(istate); + const elecstate::DensityMatrix DM_trans = this->cal_transition_density_matrix(istate); for (int is = 0;is < this->nspin_x;++is) { // 2. transition density @@ -86,7 +87,7 @@ ModuleBase::Vector3> LR::LR_Spectrum>: //1. transition density matrix ModuleBase::Vector3> trans_dipole(0.0, 0.0, 0.0); - const elecstate::DensityMatrix, std::complex>& DM_trans = this->cal_transition_density_matrix(istate); + const elecstate::DensityMatrix, std::complex> DM_trans = this->cal_transition_density_matrix(istate); for (int is = 0;is < this->nspin_x;++is) { // 2. transition density @@ -149,6 +150,7 @@ void LR::LR_Spectrum::cal_transition_dipoles_length() { transition_dipole_.resize(nstate); this->mean_squared_transition_dipole_.resize(nstate); + for (int istate = 0;istate < nstate;++istate) { transition_dipole_[istate] = cal_transition_dipole_istate_length(istate); @@ -165,9 +167,11 @@ void LR::LR_Spectrum::oscillator_strength() double osc_tot = 0.0; for (int istate = 0;istate < nstate;++istate) { - osc[istate] = this->mean_squared_transition_dipole_[istate] * this->eig[istate] * 2.; + osc[istate] = this->mean_squared_transition_dipole_[istate] * this->omega[istate] * 2.; osc_tot += osc[istate] / 2.; //Ry to Hartree (1/2) } + const int nele = (this->nspin_x == 2) ? this->nocc[0] + this->nocc[1] : 2 * this->nocc[0]; + osc_tot /= this->nk * nele; check_sum_rule(osc_tot); } @@ -200,8 +204,8 @@ void LR::LR_Spectrum::optical_absorption_method1(const std::vector& f { std::complex f_complex = std::complex(freq[f], eta); double abs = 0.0; - // for (int i = 0;i < osc.size();++i) { abs += (osc[i] / (f_complex * f_complex - eig[i] * eig[i])).imag() * freq[f] * FourPI_div_c; } - for (int i = 0;i < osc.size();++i) { abs += (osc[i] / (f_complex * f_complex - eig[i] * eig[i])).imag() * fac; } + // for (int i = 0;i < osc.size();++i) { abs += (osc[i] / (f_complex * f_complex - omega[i] * omega[i])).imag() * freq[f] * FourPI_div_c; } + for (int i = 0;i < osc.size();++i) { abs += (osc[i] / (f_complex * f_complex - omega[i] * omega[i])).imag() * fac; } if (GlobalV::MY_RANK == 0) { ofs << freq[f] * ModuleBase::Ry_to_eV << "\t" << 91.126664 / freq[f] << "\t" << std::abs(abs) << std::endl; } } ofs.close(); @@ -211,58 +215,200 @@ template void LR::LR_Spectrum::transition_analysis(const std::string& spintype) { ModuleBase::TITLE("LR::LR_Spectrum", "transition_analysis"); - std::ofstream& ofs = GlobalV::ofs_running; - ofs << "==================================================================== " << std::endl; - ofs << std::setw(40) << spintype << std::endl; - ofs << "==================================================================== " << std::endl; - ofs << std::setw(8) << "State" << std::setw(30) << "Excitation Energy (Ry, eV)" << - std::setw(90) << "Transition dipole x, y, z (a.u.)" << std::setw(30) << "Oscillator strength(a.u.)" << std::endl; - ofs << "------------------------------------------------------------------------------------ " << std::endl; - for (int istate = 0;istate < nstate;++istate) - ofs << std::setw(8) << istate << std::setw(15) << std::setprecision(6) << eig[istate] << std::setw(15) << eig[istate] * ModuleBase::Ry_to_eV - << std::setprecision(4) << std::setw(30) << transition_dipole_[istate].x << std::setw(30) << transition_dipole_[istate].y << std::setw(30) << transition_dipole_[istate].z - << std::setprecision(6) << std::setw(30) << oscillator_strength_[istate] << std::endl; - ofs << "------------------------------------------------------------------------------------ " << std::endl; - ofs << std::setw(8) << "State" << std::setw(20) << "Occupied orbital" - << std::setw(20) << "Virtual orbital" << std::setw(30) << "Excitation amplitude" - << std::setw(30) << "Excitation rate" - << std::setw(10) << "k-point" << std::endl; - ofs << "------------------------------------------------------------------------------------ " << std::endl; - for (int istate = 0;istate < nstate;++istate) + ModuleBase::timer::start("LR_Spectrum", "transition_analysis"); + std::ofstream ofs; + std::ofstream ofs_k; + const int nbands = nocc[0] + nvirt[0]; + const bool use_td_weight = (this->vmo_ptr != nullptr && LR_Util::tolower(this->gauge) == "velocity"); + if (GlobalV::MY_RANK == 0) + { + ofs.open(PARAM.globalv.global_out_dir + "trans_analysis_" + spintype + ".dat"); + ofs_k.open(PARAM.globalv.global_out_dir + "trans_kweight_" + spintype + ".dat"); + ofs << "==================================================================== \n"; + ofs << std::setw(40) << spintype << '\n'; + ofs << "==================================================================== \n"; + ofs << std::setw(8) << "State" << std::setw(30) << "Excitation Energy (Ry, eV)" << + std::setw(90) << "Transition dipole x, y, z (a.u.)" << std::setw(30) << "Oscillator strength(a.u.)" << '\n'; + ofs << "------------------------------------------------------------------------------------ \n"; + for (int istate = 0;istate < nstate;++istate) + ofs << std::setw(8) << istate << std::setw(15) << std::setprecision(6) << omega[istate] + << std::setw(15) << omega[istate] * ModuleBase::Ry_to_eV + << std::setprecision(4) << std::setw(30) << transition_dipole_[istate].x + << std::setw(30) << transition_dipole_[istate].y << std::setw(30) << transition_dipole_[istate].z + << std::setprecision(6) << std::setw(30) << oscillator_strength_[istate] << std::endl; + ofs << "------------------------------------------------------------------------------------ " << std::endl; + ofs << std::setw(8) << "State" << std::setw(20) << "Occupied orbital" + << std::setw(20) << "Virtual orbital" << std::setw(30) << "Excitation amplitude" + << std::setw(30) << "Excitation rate" + << std::setw(10) << "k-point" << '\n'; + ofs << "------------------------------------------------------------------------------------ \n"; + + ofs_k << "# Sum of exciton contribution of k-point.\n"; + ofs_k << "# weight1(k) = sum_{state,spin,occ,virt} |X(state,spin,k,occ,virt)|^2+|Y(state,spin,k,occ,virt)|^2.\n"; + ofs_k << "# weight2(k) = sum_{state,spin,direction,occ,virt} |td * X|^2 + |td *Y|^2, where td index as (spin,dir,k,occ,virt).\n"; + if (!use_td_weight) + { + ofs_k << "# NOTE: td-weighted statistics require abs_gauge velocity and a valid velocity_mo pointer.\n"; + } + else { ofs_k << "# \n"; } + + ofs_k << "k-point" << std::setw(10) << "kx" << std::setw(12) << "ky" << std::setw(12) << "kz" + << std::setw(12) << "weight1" << std::setw(12) << "weight2" << '\n'; + } + std::vector local_k_weight(nk, 0.0); + std::vector local_k_td_weight(nk, 0.0); + T amp_X(0.0), amp_Y(0.0); + // Communicate only amplitudes that will be written, in batches of NCOMM states. + constexpr int NCOMM = 256; + for (int istart = 0; istart < nstate; istart += NCOMM) { - /// find the main contributions (> 0.5) - const int loffset_b = istate * ldim; - std::vector X_full(gdim, T(0));// one-band, global - for (int is = 0;is < nspin_x;++is) + const int iend = std::min(istart + NCOMM, nstate); + const std::size_t ncount_c = std::size_t(iend - istart) * this->gdim; + if (ncount_c > std::numeric_limits::max()) { - const int loffset_bs = loffset_b + is * nk * pX[0].get_local_size(); - const int goffset_s = is * nk * nocc[0] * nvirt[0]; - for (int ik = 0;ik < nk;++ik) + throw std::overflow_error("in transition_analysis: overflow converting to int!"); + } + + std::vector local_indices; + std::vector local_amplitudes; + for (int istate = istart; istate < iend; ++istate) + { + const int loffset_b = istate * ldim; + const int goffset_b = (istate - istart) * gdim; + for (int is = 0;is < nspin_x;++is) { - const int loffset_x = loffset_bs + ik * pX[is].get_local_size(); - const int goffset_x = goffset_s + ik * nocc[is] * nvirt[is]; -#ifdef __MPI - LR_Util::gather_2d_to_full(this->pX[is], X + loffset_x, X_full.data() + goffset_x, false, nvirt[is], nocc[is]); -#endif + const int loffset_bs = loffset_b + is * nk * pX[0].get_local_size(); + const int goffset_bs = goffset_b + is * nk * nocc[0] * nvirt[0]; + const int goffset_v_s = is * 3 * nk * nbands * nbands; + const int eig_ks_spin_offset = is * nk * nbands; + for (int ik = 0;ik < nk;++ik) + { + const int loffset_x = loffset_bs + ik * pX[is].get_local_size(); + const int goffset_x = goffset_bs + ik * nocc[is] * nvirt[is]; + const int goffset_v_k = goffset_v_s + ik * nbands * nbands; + const int eig_ks_k_offset = eig_ks_spin_offset + ik * nbands; + for (int io = 0; io < pX[is].get_col_size(); ++io) + { + const int iocc = pX[is].local2global_col(io); + for (int iv = 0; iv < pX[is].get_row_size(); ++iv) + { + const int ivirt = pX[is].local2global_row(iv); + amp_X = X[loffset_x + io * pX[is].get_row_size() + iv]; + if (this->is_full && this->Y ) { + amp_Y = Y[loffset_x + io * pX[is].get_row_size() + iv]; + } + local_k_weight[ik] += std::norm(amp_X) + std::norm(amp_Y); + if (use_td_weight) + { + const int goffset_v = goffset_v_k + (ivirt + nocc[is]) * nbands + iocc; + const double gap = (eig_ks[eig_ks_k_offset + nocc[is] + ivirt] + - eig_ks[eig_ks_k_offset + iocc]) / ModuleBase::e2; // Ry to Hartree + for (int id = 0; id < 3; ++id) + { + const int v_index = goffset_v + id * nk * nbands * nbands;; + std::complex td_X = ModuleBase::IMAG_UNIT * this->vmo_ptr[v_index] * amp_X / gap; + std::complex td_Y = -1.0 * ModuleBase::IMAG_UNIT * std::conj(this->vmo_ptr[v_index]) * amp_Y / gap; + // |X_{aik}/(Ea-Ei)|^2 + |Y_{aik}/(Ei-Ea)|^2 + if (this->nspin_x == 1) { td_X *= sqrt(2.0); td_Y *= sqrt(2.0); } + local_k_td_weight[ik] += std::norm(td_X) + std::norm(td_Y); + } + } + if (std::abs(amp_X) > ana_thr) // only X components temporarily + { + local_indices.push_back(goffset_x + iocc * nvirt[is] + ivirt); + local_amplitudes.push_back(amp_X); + } + } + } + } } } - std::map> abs_order; - for (int i = 0;i < gdim;++i) { double abs = std::abs(X_full.at(i));if (abs > ana_thr) { abs_order[abs] = i; } } - if (abs_order.size() > 0) { - for (auto it = abs_order.cbegin();it != abs_order.cend();++it) + + std::vector all_indices; + std::vector all_amplitudes; + #ifdef __MPI + const int local_count = static_cast(local_indices.size()); + int comm_size = 0; + MPI_Comm_size(this->pX[0].comm(), &comm_size); + std::vector recv_counts; + std::vector displs; + if (GlobalV::MY_RANK == 0) + { + recv_counts.resize(comm_size); + } + MPI_Gather(&local_count, 1, MPI_INT, recv_counts.data(), 1, MPI_INT, 0, this->pX[0].comm()); + if (GlobalV::MY_RANK == 0) + { + displs.resize(comm_size, 0); + for (int ip = 1; ip < comm_size; ++ip) { - auto pair_info = get_pair_info(it->second); + displs[ip] = displs[ip - 1] + recv_counts[ip - 1]; + } + const std::size_t total_count = std::size_t(displs.back()) + recv_counts.back(); + all_indices.resize(total_count); + all_amplitudes.resize(total_count); + } + MPI_Gatherv(local_indices.data(), local_count, MPI_INT, + all_indices.data(), recv_counts.data(), displs.data(), MPI_INT, + 0, this->pX[0].comm()); + MPI_Gatherv(local_amplitudes.data(), local_count, LR_Util::MPIType::value(), + all_amplitudes.data(), recv_counts.data(), displs.data(), LR_Util::MPIType::value(), + 0, this->pX[0].comm()); + #else + all_indices = std::move(local_indices); + all_amplitudes = std::move(local_amplitudes); + #endif + if (GlobalV::MY_RANK != 0) continue; // only rank 0 write the analysis file + + std::vector>> contributions(iend - istart); + for (std::size_t i = 0; i < all_indices.size(); ++i) + { + const int ibatch = all_indices[i] / gdim; + contributions[ibatch].emplace_back(all_indices[i] - ibatch * gdim, all_amplitudes[i]); + } + + for (int istate = istart; istate < iend; ++istate) + { + auto& state_contributions = contributions[istate - istart]; + std::sort(state_contributions.begin(), state_contributions.end(), + [](const std::pair& l, const std::pair& r) { return std::abs(l.second) > std::abs(r.second); }); + + for (auto it = state_contributions.cbegin(); it != state_contributions.cend(); ++it) + { + auto pair_info = get_pair_info(it->first); const int& is = pair_info["ispin"]; const std::string s = nspin_x == 2 ? (is == 0 ? "a" : "b") : ""; - ofs << std::setw(8) << (it == abs_order.cbegin() ? std::to_string(istate) : " ") + ofs << std::setw(8) << (it == state_contributions.cbegin() ? std::to_string(istate) : " ") << std::setw(20) << std::to_string(pair_info["iocc"] + 1) + s << std::setw(20) << std::to_string(pair_info["ivirt"] + nocc[is] + 1) + s// iocc and ivirt - << std::setw(30) << X_full.at(it->second) - << std::setw(30) << std::norm(X_full.at(it->second)) - << std::setw(10) << pair_info["ik"] + 1 << std::endl; + << std::setw(30) << it->second + << std::setw(30) << std::norm(it->second) + << std::setw(10) << pair_info["ik"] + 1 << '\n'; } } } - ofs << "==================================================================== " << std::endl; + std::vector k_weight(nk, 0.0); + std::vector k_td_weight(nk, 0.0); +#ifdef __MPI + MPI_Reduce(local_k_weight.data(), k_weight.data(), nk, MPI_DOUBLE, MPI_SUM, 0, this->pX[0].comm()); + MPI_Reduce(local_k_td_weight.data(), k_td_weight.data(), nk, MPI_DOUBLE, MPI_SUM, 0, this->pX[0].comm()); +#else + k_weight = std::move(local_k_weight); + k_td_weight = std::move(local_k_td_weight); +#endif + if (GlobalV::MY_RANK == 0) + { + ofs_k << std::fixed << std::setprecision(5); + for (int ik = 0; ik < nk; ++ik) + { + ofs_k << std::setw(5) << ik + 1 << std::setw(12) << kv.kvec_d[ik].x + << std::setw(12) << kv.kvec_d[ik].y << std::setw(12) << kv.kvec_d[ik].z + << std::setw(12) << k_weight[ik] << std::setw(12) << k_td_weight[ik] << '\n'; + } + ofs.close(); + ofs_k.close(); + } + ModuleBase::timer::end("LR_Spectrum", "transition_analysis"); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "LR::LR_Spectrum::transition_analysis"); } template @@ -283,15 +429,17 @@ void LR::LR_Spectrum::write_transition_dipole(const std::string& filename) { std::ofstream ofs(filename); ofs << "Transition dipole moment (a.u.)" << std::endl; - ofs << std::setw(20) << "State" << std::setw(20) << "x" << std::setw(20) << "y" << std::setw(20) << "z" << std::setw(20) << "average" << std::endl; + ofs << std::setw(6) << "State" << std::setw(13) << "Energy (eV)" << std::setw(15) << "x" << std::setw(23) << "|x|^2" << std::setw(19) << "y" << std::setw(23) <<"|y|^2" << std::setw(19) << "z" << std::setw(23) <<"|z|^2" << std::setw(13) << "average" << std::endl; for (int istate = 0;istate < nstate;++istate) { - ofs << std::setw(20) << istate << std::setw(20) << transition_dipole_[istate].x << std::setw(20) - << transition_dipole_[istate].y << std::setw(20) - << transition_dipole_[istate].z << std::setw(20) - << mean_squared_transition_dipole_[istate] << std::endl; + ofs << std::setw(4) << istate << std::setw(13) << std::setprecision(6) << omega[istate] * ModuleBase::Ry_to_eV + << std::setw(29) << transition_dipole_[istate].x << std::setw(13) << std::norm(transition_dipole_[istate].x) + << std::setw(29) << transition_dipole_[istate].y << std::setw(13) << std::norm(transition_dipole_[istate].y) + << std::setw(29) << transition_dipole_[istate].z << std::setw(13) << std::norm(transition_dipole_[istate].z) + << std::setw(13) << mean_squared_transition_dipole_[istate] << std::endl; } ofs.close(); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "LR::LR_Spectrum::write_transition_dipole " + filename); } template class LR::LR_Spectrum; diff --git a/source/source_lcao/module_lr/lr_spectrum.h b/source/source_lcao/module_lr/lr_spectrum.h index 5ed4678dc9..57723cccf4 100644 --- a/source/source_lcao/module_lr/lr_spectrum.h +++ b/source/source_lcao/module_lr/lr_spectrum.h @@ -5,6 +5,8 @@ #include "source_lcao/module_lr/utils/lr_util.h" #include "source_basis/module_nao/two_center_bundle.h" #include "source_lcao/module_rt/velocity_op.h" +#include "source_lcao/module_lr/utils/spectrum_mo.hpp" + namespace LR { template @@ -16,20 +18,27 @@ namespace LR const UnitCell& ucell, const K_Vectors& kv_in, const Grid_Driver& gd, const std::vector& orb_cutoff, const TwoCenterBundle& two_center_bundle_, const std::vector& pX_in, const Parallel_2D& pc_in, const Parallel_Orbitals& pmat_in, - const double* eig, const T* X, const int& nstate, const bool& openshell, - const std::string& gauge = "length") : - nspin_x(openshell ? 2 : 1), naos(naos), nocc(nocc), nvirt(nvirt), nk(kv_in.get_nks() / nspin_global), + const double* omega, const double* eig_ks, const T* X, const int& nstate, const bool& openshell, + const std::string& gauge) : + nspin_x(openshell ? 2 : 1), naos(naos), nocc(nocc), nvirt(nvirt), + nk(nspin_global == 2 ? kv_in.get_nks() / 2 : kv_in.get_nks()), rho_basis(rho_basis), ucell(ucell), kv(kv_in), gd_(gd), orb_cutoff_(orb_cutoff), two_center_bundle_(two_center_bundle_), pX(pX_in), pc(pc_in), pmat(pmat_in), - eig(eig), X(X), nstate(nstate), + omega(omega), eig_ks(eig_ks), X(X), nstate(nstate), gauge(gauge), ldim(nk* (nspin_x == 2 ? pX_in[0].get_local_size() + pX_in[1].get_local_size() : pX_in[0].get_local_size())), gdim(nk* std::inner_product(nocc.begin(), nocc.end(), nvirt.begin(), 0)) { - for (int is = 0;is < nspin_global;++is) { psi_ks.emplace_back(LR_Util::get_psi_spin(psi_ks_in, is, nk)); } - gauge == "velocity" ? this->cal_transition_dipoles_velocity() : this->cal_transition_dipoles_length(); - this->oscillator_strength(); + for (int is = 0;is < nspin_global;++is) { psi_ks_vec.emplace_back(LR_Util::get_psi_spin(psi_ks_in, is, nk)); } }; + void set_vmo(std::complex* vmo_in) { this->vmo_ptr = vmo_in; }; + void set_Y(T* Y_in) { this->Y = Y_in; }; + void set_full(bool tag) { this->is_full = tag; }; + void cal_spectrum(){ + std::cout << "Calculating transition dipole moments in " << this->gauge << " gauge." << std::endl; + this->gauge == "velocity" ? this->cal_transition_dipoles_velocity(this->eig_ks) : this->cal_transition_dipoles_length(); + this->oscillator_strength(); + } /// @brief calculate the optical absorption spectrum with $Im[1/[(w+i\eta)^2-\Omega_S^2]]$ void optical_absorption_method1(const std::vector& freq, const double eta); /// @brief calculate the optical absorption spectrum with lorentzian delta function @@ -40,10 +49,15 @@ namespace LR //========================================== test functions ============================================== /// @brief write transition dipole void write_transition_dipole(const std::string& filename); - /// @brief calculate transition dipole in velocity gauge using ks eigenvalues instead of excitation energies - void test_transition_dipoles_velocity_ks(const double* const ks_eig); + /// @brief calculate transition dipole in velocity gauge using excitation eigenvalues instead of KS differences + void test_transition_dipoles_velocity_omega(); + //====================================================================================================== private: + enum class DipoleEnergyType { + LR_EIG, + KS_GAP + }; /// $$2/3\Omega\sum_{ia\sigma} |\braket{\psi_{i}|\mathbf{r}|\psi_{a}} |^2\int \rho_{\alpha\beta}(\mathbf{r}) \mathbf{r} d\mathbf{r}$$ void oscillator_strength(); /// calculate the transition dipole of state S in length gauge: $\sum_{iak}X^S_{iak}$ @@ -53,11 +67,13 @@ namespace LR /// calculate the transition dipole of state S in velocity gauge: $i(\sum_{iak}X^S_{iak})/\Omega_S$ ModuleBase::Vector3 cal_transition_dipole_istate_velocity_R(const int istate, const Velocity_op>& vR); ModuleBase::Vector3 cal_transition_dipole_istate_velocity_k(const int istate, const Velocity_op>& vR); + void cal_transition_dipole_istate_velocity_mo(DipoleEnergyType method, const std::vector& eig_ks_diff); /// calculate the transition dipole of all states in velocity gauge - void cal_transition_dipoles_velocity(); + void cal_transition_dipoles_velocity(const double* const eig_ks); double cal_mean_squared_dipole(ModuleBase::Vector3 dipole); /// calculate the transition density matrix elecstate::DensityMatrix cal_transition_density_matrix(const int istate, const T* X_in = nullptr, const bool need_R = true); + const int nspin_x = 1; ///< 1 for singlet/triplet, 2 for updown(openshell) const int naos = 1; const std::vector& nocc; @@ -67,10 +83,13 @@ namespace LR const int ldim = 1;///< local leading dimension of X, or the data size of each state const int gdim = 1;///< global leading dimension of X const double ana_thr = 0.3; ///< {abs(X) > thr} will appear in the transition analysis log - const double* eig = nullptr; - const T* X = nullptr; + const double* omega; ///< excitation energies + const double* eig_ks; ///< KS eigenvalues + const T* X; + T* Y = nullptr; ///< the deexcitation part of amplitudes + bool is_full = false; const K_Vectors& kv; - std::vector> psi_ks; + std::vector> psi_ks_vec; const std::vector& pX; const Parallel_2D& pc; const Parallel_Orbitals& pmat; @@ -79,10 +98,12 @@ namespace LR const UnitCell& ucell; const std::vector& orb_cutoff_; const TwoCenterBundle& two_center_bundle_; + const std::string gauge; void cal_gint_rho(double** rho, const int& nrxx); std::map get_pair_info(const int i); ///< given the index in X, return its ispin, ik, iocc, ivirt + std::complex* vmo_ptr = nullptr; ///< pointer to velocity matrix elements in MO basis, used in velocity gauge std::vector> transition_dipole_; ///< $\braket{ \psi_{i} | \mathbf{r} | \psi_{a} }$ std::vector mean_squared_transition_dipole_; /// $|dipole|^2/3$, atomic unit (Hartree) std::vector oscillator_strength_;///< $2/3\Omega |\sum_{ia\sigma} \braket{\psi_{i}|\mathbf{r}|\psi_{a}} |^2$, atomic unit (Hartree) diff --git a/source/source_lcao/module_lr/lr_spectrum_velocity.cpp b/source/source_lcao/module_lr/lr_spectrum_velocity.cpp index c3259da039..60573cbbc1 100644 --- a/source/source_lcao/module_lr/lr_spectrum_velocity.cpp +++ b/source/source_lcao/module_lr/lr_spectrum_velocity.cpp @@ -19,8 +19,8 @@ namespace LR inp.out_element_info, inp.cal_force); // actually this class calculates the velocity matrix v(R) at A=0 Velocity_op> vR(&ucell, &gd, &pmat, orb, two_center_bundle.overlap_orb.get()); - vR.calculate_vcomm_r(); // $<\mu, 0|[Vnl, r]|\nu, R>$ - vR.calculate_grad_term(); // $<\mu, 0|\nabla|\nu, R>$ + vR.calculate_vcomm_r(); // $<\mu, 0|i[Vnl, r]|\nu, R>$ + vR.calculate_grad_term(); // $<\mu, 0|-i∇r|\nu, R>$ return vR; } @@ -33,6 +33,25 @@ namespace LR const double c = eta_au / std::sqrt(2. * std::log(2.)); return std::exp(-dfreq_au * dfreq_au / (2 * c * c)) / (std::sqrt(2 * M_PI) * c); } + template + void LR::LR_Spectrum::optical_absorption_method2(const std::vector& freq, const double eta) + { + ModuleBase::TITLE("LR::LR_Spectrum", "optical_absorption_method2"); + // 4*pi^2/V * mean_squared_dipole *delta(w-Omega_S) + std::ofstream ofs(PARAM.globalv.global_out_dir + "absorption.dat"); + if (GlobalV::MY_RANK == 0) { ofs << "Frequency (eV) | wave length(nm) | Absorption (a.u.)" << std::endl; } + const double fac = 4 * M_PI * M_PI / ucell.omega / this->nk; + for (int f = 0;f < freq.size();++f) + { + double abs_value = 0.0; + for (int i = 0;i < nstate;++i) + { + abs_value += this->mean_squared_transition_dipole_[i] * lorentz_delta((freq[f] - omega[i]) / ModuleBase::e2, eta / ModuleBase::e2); // e2: Ry to Hartree + } + abs_value *= fac; + if (GlobalV::MY_RANK == 0) { ofs << freq[f] * ModuleBase::Ry_to_eV << "\t" << 91.126664 / freq[f] << "\t" << abs_value << std::endl; } + } + } template inline ModuleBase::Vector3 convert_vector_to_vector3(const std::vector>& vec); template<> inline ModuleBase::Vector3 convert_vector_to_vector3(const std::vector>& vec) @@ -46,15 +65,28 @@ namespace LR return ModuleBase::Vector3>(vec[0], vec[1], vec[2]); } - /// this algorithm has bug in multi-k cases, just for test + template + inline ModuleBase::Vector3 convert_ptr_to_vector3(const std::complex* ptr); + template<> + inline ModuleBase::Vector3 convert_ptr_to_vector3(const std::complex* ptr) + { + return ModuleBase::Vector3(ptr[0].real(), ptr[1].real(), ptr[2].real()); + } + template<> + inline ModuleBase::Vector3> convert_ptr_to_vector3>(const std::complex* ptr) + { + return ModuleBase::Vector3>(ptr[0], ptr[1], ptr[2]); + } + + /// this algorithm is stored for reference template ModuleBase::Vector3 LR::LR_Spectrum::cal_transition_dipole_istate_velocity_R(const int istate, const Velocity_op>& vR) { // transition density matrix D(R) const elecstate::DensityMatrix& DM_trans = this->cal_transition_density_matrix(istate); - std::vector> trans_dipole(3, 0.0); // $=\sum_{uvR} v(R) D(R) = \sum_{iak}X_{iak}$ - const std::complex fac = ModuleBase::IMAG_UNIT / (eig[istate] / ModuleBase::e2); // eV to Hartree + std::vector> trans_dipole(3, 0.0); // $=\sum_{uvR} v(R) D(R) = \sum_{aik}X_{aik}$ + const std::complex fac = ModuleBase::IMAG_UNIT / (omega[istate] / ModuleBase::e2); // Ry to Hartree for (int i = 0; i < 3; i++) { for (int is = 0;is < this->nspin_x; ++is) @@ -68,15 +100,15 @@ namespace LR return convert_vector_to_vector3(trans_dipole); } - // this algorithm is actually in use + // this algorithm is stored for reference template ModuleBase::Vector3 LR::LR_Spectrum::cal_transition_dipole_istate_velocity_k(const int istate, const Velocity_op>& vR) { // transition density matrix D(R) const elecstate::DensityMatrix& DM_trans = this->cal_transition_density_matrix(istate, this->X, false); - std::vector> trans_dipole(3, 0.0); // $=\sum_{uvR} v(R) D(R) = \sum_{iak}X_{iak}$ - const std::complex fac = ModuleBase::IMAG_UNIT / (eig[istate] / ModuleBase::e2); // eV to Hartree + std::vector> trans_dipole(3, 0.0); // $=\sum_{uvk} v(k) D(k) = \sum_{aik}X_{aik}$ + const std::complex fac = ModuleBase::IMAG_UNIT / (omega[istate] / ModuleBase::e2); // Ry to Hartree for (int i = 0; i < 3; i++) { for (int is = 0;is < this->nspin_x;++is) @@ -95,93 +127,156 @@ namespace LR return convert_vector_to_vector3(trans_dipole); } + // this algorithm is faster since velocity_mo is calculated and each transition state only need to contract with X template - void LR::LR_Spectrum::cal_transition_dipoles_velocity() + void LR::LR_Spectrum::cal_transition_dipole_istate_velocity_mo(DipoleEnergyType method, const std::vector& eig_ks_diff) { - const Velocity_op>& vR = get_velocity_matrix_R(ucell, gd_, pmat, two_center_bundle_); // velocity matrix v(R) - transition_dipole_.resize(nstate); - this->mean_squared_transition_dipole_.resize(nstate); - for (int istate = 0;istate < nstate;++istate) + ModuleBase::timer::start("LR_Spectrum", "cal_transition_dipole_istate_velocity_mo"); + if (this->vmo_ptr == nullptr) { - transition_dipole_[istate] = cal_transition_dipole_istate_velocity_k(istate, vR); - mean_squared_transition_dipole_[istate] = cal_mean_squared_dipole(transition_dipole_[istate]); + ModuleBase::WARNING_QUIT("LR_Spectrum", "velocity_mo is null. Please pass a valid pointer."); } - } + const int nbands = this->nocc[0] + this->nvirt[0]; + assert(nbands == this->pc.get_global_col_size()); + const bool use_ks_gap = (method == DipoleEnergyType::KS_GAP); + + std::vector> trans_dipole_buf(3 * nstate, 0.0); // $= \sum_{aik} i (X_{aik}/(Ea-Ei) + Y_{aik}/(Ei-Ea))$ + // vmo is global [spin, direction, kpoint, nbands, nbands], X is local [spin, kpoint, nocc_local, nvirt_local] - template - void LR::LR_Spectrum::optical_absorption_method2(const std::vector& freq, const double eta) - { - ModuleBase::TITLE("LR::LR_Spectrum", "optical_absorption_velocity"); - // 4*pi^2/V * mean_squared_dipole *delta(w-Omega_S) - std::ofstream ofs(PARAM.globalv.global_out_dir + "absorption.dat"); - if (GlobalV::MY_RANK == 0) { ofs << "Frequency (eV) | wave length(nm) | Absorption (a.u.)" << std::endl; } - const double fac = 4 * M_PI * M_PI / ucell.omega / this->nk; - for (int f = 0;f < freq.size();++f) +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int istate = 0; istate < nstate; ++istate) { - double abs_value = 0.0; - for (int i = 0;i < nstate;++i) + const std::complex fac = use_ks_gap ? ModuleBase::IMAG_UNIT : + (ModuleBase::IMAG_UNIT / (omega[istate] / 2.0)); // Ry to Hartree; + const std::size_t loffset_X_b = std::size_t(istate) * std::size_t(this->ldim); + for (int id = 0; id < 3; ++id) { - abs_value += this->mean_squared_transition_dipole_[i] * lorentz_delta((freq[f] - eig[i]) / ModuleBase::e2, eta / ModuleBase::e2); // e2: Ry to Hartree - } - abs_value *= fac; - if (GlobalV::MY_RANK == 0) { ofs << freq[f] * ModuleBase::Ry_to_eV << "\t" << 91.126664 / freq[f] << "\t" << abs_value << std::endl; } + std::complex td = 0.0; // short name of transition dipole + for (int is = 0; is < this->nspin_x; ++is) + { + const std::size_t loffset_X_bs = loffset_X_b + is * nk * pX[0].get_local_size(); + const int goffset_v_ds = (is * 3 + id) * nk * nbands * nbands; + for (int ik = 0; ik < nk; ++ik) + { + const double wk = this->kv.wk[ik] * nk; // k-point weight, normalized as sum = nk + const std::size_t loffset_X = loffset_X_bs + ik * pX[is].get_local_size(); + const int goffset_v = goffset_v_ds + ik * nbands * nbands; + for (int io = 0; io < pX[is].get_col_size(); ++io) // nocc_local + { + int io_g = pX[is].local2global_col(io); + for (int iv = 0; iv < pX[is].get_row_size(); ++iv) // nvirt_local + { + int iv_g = pX[is].local2global_row(iv); + const std::size_t X_index = loffset_X + io * pX[is].get_row_size() + iv; + const int v_index = goffset_v + (iv_g+nocc[is]) * nbands + io_g; + if (use_ks_gap) + { + td += this->vmo_ptr[v_index] * X[X_index] / eig_ks_diff[X_index - loffset_X_b] * wk; + if (this->is_full) + { + // THE HERMITIAN CONJUGATE OF VMO HAS BEEN VERIFIED + // const int v_index2 = goffset_v + io_g * nbands + iv_g + nocc[is]; + // if (std::abs(vmo_ptr[v_index2] - std::conj(this->vmo_ptr[v_index])) > 1e-5){ + // std::cout<<"io:"<X_{aik}/(Ea-Ei) + Y_{aik}/(Ei-Ea) + td -= std::conj(this->vmo_ptr[v_index]) * Y[X_index] / eig_ks_diff[X_index - loffset_X_b] * wk; + } + } + else + { + td += this->vmo_ptr[v_index] * X[X_index] * wk; + if (this->is_full) + { + td += std::conj(this->vmo_ptr[v_index]) * Y[X_index] * wk; + } + } + } + } + } + } // end for spin_x, only matter in open-shell system + td *= fac; + if (this->nspin_x == 1) { td *= sqrt(2.0); } // *2 for 2 spins, /sqrt(2) for the halfed dimension of X in the normalizaiton + trans_dipole_buf[3 * istate + id] = td; + } // end for direction + } + Parallel_Reduce::reduce_all(trans_dipole_buf.data(), 3 * nstate); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int istate = 0; istate < nstate; ++istate) + { + std::complex* ptr = &trans_dipole_buf[3 * istate]; + this->transition_dipole_[istate] = convert_ptr_to_vector3(ptr); + this->mean_squared_transition_dipole_[istate] = cal_mean_squared_dipole(transition_dipole_[istate]); } + ModuleBase::timer::end("LR_Spectrum", "cal_transition_dipole_istate_velocity_mo"); + } + + template + void LR::LR_Spectrum::test_transition_dipoles_velocity_omega() + { + ModuleBase::timer::start("LR_Spectrum", "test_transition_dipoles_velocity_omega"); + this->transition_dipole_.resize(nstate); + this->mean_squared_transition_dipole_.resize(nstate); + this->cal_transition_dipole_istate_velocity_mo(DipoleEnergyType::LR_EIG, {}); + this->oscillator_strength(); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "LR::LR_Spectrum::test_transition_dipoles_velocity_omega"); + ModuleBase::timer::end("LR_Spectrum", "test_transition_dipoles_velocity_omega"); } inline void cal_eig_ks_diff(double* const eig_ks_diff, const double* const eig_ks, const Parallel_2D& px, const int nk, const int nocc, const int nvirt) { +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif for (int ik = 0;ik < nk;++ik) { - const int& start_k = ik * (nocc + nvirt); + const int start_k = ik * (nocc + nvirt); for (int io = 0;io < px.get_col_size();++io) //nocc_local { for (int iv = 0;iv < px.get_row_size();++iv) //nvirt_local { int io_g = px.local2global_col(io); int iv_g = px.local2global_row(iv); - eig_ks_diff[ik * px.get_local_size() + io * px.get_row_size() + iv] = (eig_ks[start_k + nocc + iv_g] - eig_ks[start_k + io_g]) / ModuleBase::e2; // eV to Hartree + eig_ks_diff[ik * px.get_local_size() + io * px.get_row_size() + iv] = (eig_ks[start_k + nocc + iv_g] - eig_ks[start_k + io_g]) / ModuleBase::e2; // Ry to Hartree } } } } template - void LR::LR_Spectrum::test_transition_dipoles_velocity_ks(const double* const ks_eig) + void LR::LR_Spectrum::cal_transition_dipoles_velocity(const double* const eig_ks) { - // velocity matrix v(R) - const Velocity_op>& vR = get_velocity_matrix_R(ucell, gd_, pmat, two_center_bundle_); + ModuleBase::timer::start("LR_Spectrum", "cal_transition_dipoles_velocity"); + + this->transition_dipole_.resize(nstate); + this->mean_squared_transition_dipole_.resize(nstate); + // (e_c-e_v) of KS eigenvalues std::vector eig_ks_diff(this->ldim); for (int is = 0;is < this->nspin_x;++is) { - cal_eig_ks_diff(eig_ks_diff.data() + is * nk * pX[0].get_local_size(), ks_eig, pX[is], nk, nocc[is], nvirt[is]); - } - // X/(ec-ev) - std::vector X_div_ks_eig(nstate * this->ldim); - for (int istate = 0;istate < nstate;++istate) - { - const int st = istate * this->ldim; - std::transform(X + st, X + st + ldim, eig_ks_diff.begin(), X_div_ks_eig.data() + st, std::divides()); + cal_eig_ks_diff(eig_ks_diff.data() + is * nk * pX[0].get_local_size(), eig_ks, pX[is], nk, nocc[is], nvirt[is]); } - this->transition_dipole_.resize(nstate); - this->mean_squared_transition_dipole_.resize(nstate); - for (int istate = 0;istate < nstate;++istate) - { - // transition density matrix D(R) - const elecstate::DensityMatrix& DM_trans = this->cal_transition_density_matrix(istate, X_div_ks_eig.data()); - std::vector> tmp_trans_dipole(3, 0.0); - for (int i = 0; i < 3; i++) - { - for (int is = 0;is < this->nspin_x; ++is) - { - tmp_trans_dipole[i] += LR_Util::dot_R_matrix(*vR.get_current_term_pointer(i), *DM_trans.get_DMR_pointer(is + 1), ucell.nat) * ModuleBase::IMAG_UNIT; - } // end for spin_x, only matter in open-shell system - } // end for direction - this->transition_dipole_[istate] = convert_vector_to_vector3(tmp_trans_dipole); - this->mean_squared_transition_dipole_[istate] = cal_mean_squared_dipole(transition_dipole_[istate]); - } + // X/(ec-ev) + // std::vector X_div_ks_eig(nstate * this->ldim); + // for (int istate = 0;istate < nstate;++istate) + // { + // const int st = istate * this->ldim; + // std::transform(X + st, X + st + ldim, eig_ks_diff.begin(), X_div_ks_eig.data() + st, std::divides()); + // } + + this->cal_transition_dipole_istate_velocity_mo(DipoleEnergyType::KS_GAP, eig_ks_diff); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "LR::LR_Spectrum::cal_transition_dipoles_velocity"); + ModuleBase::timer::end("LR_Spectrum", "cal_transition_dipoles_velocity"); } -} + +} // namespace LR + template class LR::LR_Spectrum; template class LR::LR_Spectrum>; \ No newline at end of file diff --git a/source/source_lcao/module_lr/operator_casida/operator_lr_diag.h b/source/source_lcao/module_lr/operator_casida/operator_lr_diag.h index e29ca438b8..25ea986f1e 100644 --- a/source/source_lcao/module_lr/operator_casida/operator_lr_diag.h +++ b/source/source_lcao/module_lr/operator_casida/operator_lr_diag.h @@ -44,10 +44,12 @@ namespace LR const bool is_first_node = false)const override { ModuleBase::TITLE("OperatorLRDiag", "act"); + ModuleBase::timer::start("OperatorLRDiag", "act"); ModuleBase::vector_mul_vector_op()(nk * pX.get_local_size(), // local size of particle-hole basis hpsi, psi_in, this->eig_ks_diff.c); + ModuleBase::timer::end("OperatorLRDiag", "act"); } private: const Parallel_2D& pX; diff --git a/source/source_lcao/module_lr/operator_casida/operator_lr_exx.cpp b/source/source_lcao/module_lr/operator_casida/operator_lr_exx.cpp index 8ba3a946e2..3879569b2f 100644 --- a/source/source_lcao/module_lr/operator_casida/operator_lr_exx.cpp +++ b/source/source_lcao/module_lr/operator_casida/operator_lr_exx.cpp @@ -15,9 +15,8 @@ namespace LR for (int iat2 = 0;iat2 < ucell.nat;++iat2) { const int it2 = ucell.iat2it[iat2]; for (auto cell : this->BvK_cells) { - this->Ds_onebase[iat1][std::make_pair(iat2, cell)] = aims_nbasis.empty() ? - RI::Tensor({ static_cast(ucell.atoms[it1].nw), static_cast(ucell.atoms[it2].nw) }) : - RI::Tensor({ static_cast(aims_nbasis[it1]), static_cast(aims_nbasis[it2]) }); + this->Ds_onebase[iat1][std::make_pair(iat2, cell)] = + RI::Tensor({ static_cast(ucell.atoms[it1].nw), static_cast(ucell.atoms[it2].nw) }); } } } @@ -41,8 +40,8 @@ namespace LR int iat1 = ucell.itia2iat(it1, ia1); int iat2 = ucell.itia2iat(it2, ia2); auto& D2d = this->Ds_onebase[iat1][std::make_pair(iat2, cell)]; - const int nw1 = aims_nbasis.empty() ? ucell.atoms[it1].nw : aims_nbasis[it1]; - const int nw2 = aims_nbasis.empty() ? ucell.atoms[it2].nw : aims_nbasis[it2]; + const int nw1 = ucell.atoms[it1].nw; + const int nw2 = ucell.atoms[it2].nw; for (int iw1 = 0;iw1 < nw1;++iw1) for (int iw2 = 0;iw2 < nw2;++iw2) { @@ -74,8 +73,8 @@ namespace LR int iat1 = ucell.itia2iat(it1, ia1); int iat2 = ucell.itia2iat(it2, ia2); auto& D2d = this->Ds_onebase[iat1][std::make_pair(iat2, cell)]; - const int nw1 = aims_nbasis.empty() ? ucell.atoms[it1].nw : aims_nbasis[it1]; - const int nw2 = aims_nbasis.empty() ? ucell.atoms[it2].nw : aims_nbasis[it2]; + const int nw1 = ucell.atoms[it1].nw; + const int nw2 = ucell.atoms[it2].nw; for (int iw1 = 0;iw1 < nw1;++iw1) for (int iw2 = 0;iw2 < nw2;++iw2) { @@ -98,6 +97,8 @@ namespace LR const bool is_first_node)const { ModuleBase::TITLE("OperatorLREXX", "act"); + ModuleBase::timer::start("OperatorLREXX", "act"); + // convert parallel info to LibRI interfaces std::vector, std::set>> judge = RI_2D_Comm::get_2D_judge(ucell,this->pmat); @@ -112,42 +113,47 @@ namespace LR for (int ik = 0;ik < nk;++ik) { DMk_trans_pointer[ik] = &DMk_trans_vector[ik]; } // if multi-k, DM_trans(TR=double) -> Ds_trans(TR=T=complex) std::vector>>> Ds_trans = - aims_nbasis.empty() ? - RI_2D_Comm::split_m2D_ktoR(ucell,this->kv, DMk_trans_pointer, this->pmat, 1) - : RI_Benchmark::split_Ds(DMk_trans_vector, aims_nbasis, ucell); //0.5 will be multiplied + // aims_nbasis.empty() ? // ucell.nw is updated, abandoned 25-05-23 + RI_2D_Comm::split_m2D_ktoR(ucell,this->kv, DMk_trans_pointer, this->pmat, 1); + //: RI_Benchmark::split_Ds(DMk_trans_vector, aims_nbasis, ucell); //0.5 will be multiplied // LR_Util::print_CV(Ds_trans[0], "Ds_trans in OperatorLREXX", 1e-10); + // 2. cal_Hs + ModuleBase::timer::start("OperatorLREXX", "cal_Hs"); auto lri = this->exx_lri.lock(); - - // LR_Util::print_CV(Ds_trans[is], "Ds_trans in OperatorLREXX", 1e-10); lri->exx_lri.set_Ds(std::move(Ds_trans[0]), lri->info.dm_threshold); lri->exx_lri.cal_Hs(); lri->Hexxs[0] = RI::Communicate_Tensors_Map_Judge::comm_map2_first( lri->mpi_comm, std::move(lri->exx_lri.Hs), std::get<0>(judge[0]), std::get<1>(judge[0])); lri->post_process_Hexx(lri->Hexxs[0]); + // LR_Util::print_CV(lri->Hexxs[0], "Hexxs in OperatorLREXX", 1e-10); + ModuleBase::timer::end("OperatorLREXX", "cal_Hs"); // 3. set [AX]_iak = DM_onbase * Hexxs for each occ-virt pair and each k-point // caution: parrallel - - for (int io = 0;io < this->nocc;++io) + ModuleBase::timer::start("OperatorLREXX", "cal_energy"); + for (int ik = 0;ik < nk;++ik) { - for (int iv = 0;iv < this->nvirt;++iv) + for (int io = 0;io < this->nocc;++io) { - for (int ik = 0;ik < nk;++ik) + for (int iv = 0;iv < this->nvirt;++iv) { const int xstart_bk = ik * pX.get_local_size(); this->cal_DM_onebase(io, iv, ik); //set Ds_onebase for all e-h pairs (not only on this processor) // LR_Util::print_CV(Ds_onebase, "Ds_onebase of occ " + std::to_string(io) + ", virtual " + std::to_string(iv) + " in OperatorLREXX", 1e-10); - const T& ene = 2 * alpha * //minus for exchange(but here plus is right, why?), 2 for Hartree to Ry + const T& ene = 2 * alpha * //minus for exchange(but here plus, since `post_process_Hexx` has taken minus), 2 for Hartree to Ry lri->exx_lri.post_2D.cal_energy(this->Ds_onebase, lri->Hexxs[0]); if (this->pX.in_this_processor(iv, io)) { hpsi[xstart_bk + this->pX.global2local_col(io) * this->pX.get_row_size() + this->pX.global2local_row(iv)] += ene; } + //for debug + GlobalV::ofs_running << "Direct term: ik="<; template class OperatorLREXX>; diff --git a/source/source_lcao/module_lr/operator_casida/operator_lr_exx.h b/source/source_lcao/module_lr/operator_casida/operator_lr_exx.h index 2c1136201a..2e886fb4a2 100644 --- a/source/source_lcao/module_lr/operator_casida/operator_lr_exx.h +++ b/source/source_lcao/module_lr/operator_casida/operator_lr_exx.h @@ -7,7 +7,7 @@ namespace LR { - /// @brief Hxc part of A operator + /// @brief Exx part of A operator template class OperatorLREXX : public hamilt::Operator { @@ -30,19 +30,19 @@ namespace LR const Parallel_2D& pX_in, const Parallel_2D& pc_in, const Parallel_Orbitals& pmat_in, - const double& alpha = 1.0, - const std::vector& aims_nbasis = {}) + const double& alpha = 1.0) : nspin(nspin), naos(naos), nocc(nocc), nvirt(nvirt), nk(kv_in.get_nks() / nspin), psi_ks(psi_ks_in), DM_trans(DM_trans_in), exx_lri(exx_lri_in), kv(kv_in), - pX(pX_in), pc(pc_in), pmat(pmat_in), ucell(ucell_in), alpha(alpha), - aims_nbasis(aims_nbasis) + pX(pX_in), pc(pc_in), pmat(pmat_in), ucell(ucell_in), alpha(alpha) { ModuleBase::TITLE("OperatorLREXX", "OperatorLREXX"); + std::cout<<"Initializing OperatorLREXX"<cal_type = hamilt::calculation_type::lcao_exx; this->is_first_node = false; // reduce psi_ks for later use this->psi_ks_full.resize(this->nk, nocc + nvirt, this->naos); + this->psi_ks_full.zero_out(); for (int ik = 0;ik < nk;++ik) { LR_Util::gather_2d_to_full(this->pc, &this->psi_ks(ik, 0, 0), &this->psi_ks_full(ik, 0, 0), false, this->naos, nocc + nvirt); @@ -80,13 +80,12 @@ namespace LR /// ground state wavefunction const psi::Psi& psi_ks = nullptr; psi::Psi psi_ks_full; - const std::vector aims_nbasis={}; ///< number of basis functions for each type of atom in FHI-aims /// transition density matrix std::unique_ptr>& DM_trans; /// density matrix of a certain (i, a, k), with full naos*naos size for each key - /// D^{iak}_{\mu\nu}(k): 1/N_k * c^*_{ak,\mu} c_{ik,\nu} + /// D^{iak}_{\mu\nu}(k): 1/N_k * c_{ak,\mu} c^*_{ik,\nu} /// D^{iak}_{\mu\nu}(R): D^{iak}_{\mu\nu}(k)e^{-ikR} // elecstate::DensityMatrix* DM_onebase; mutable std::map>> Ds_onebase; @@ -110,7 +109,6 @@ namespace LR const Parallel_2D& pX; const Parallel_Orbitals& pmat; - // allocate Ds_onebase void allocate_Ds_onebase(); diff --git a/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.cpp b/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.cpp index cae13467a8..675a15b91d 100644 --- a/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.cpp +++ b/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.cpp @@ -19,6 +19,8 @@ namespace LR void OperatorLRHxc::act(const int nbands, const int nbasis, const int npol, const T* psi_in, T* hpsi, const int ngk_ik, const bool is_first_node)const { ModuleBase::TITLE("OperatorLRHxc", "act"); + ModuleBase::timer::start("OperatorLRHxc", "act"); + const int& sl = ispin_ks[0]; const auto psil_ks = LR_Util::get_psi_spin(psi_ks, sl, nk); @@ -45,6 +47,11 @@ namespace LR #else ao_to_mo_blas(v_hxc_2d, psil_ks, nocc[sl], nvirt[sl], hpsi); #endif + // for debug + //std::cout << "After Hxc, hpsi: [nvirt= " << nvirt[sl] << " nocc= " << nocc[sl] << " nk= " << nk << " ]" << std::endl; + //LR_Util::print_value(hpsi, nk, nocc[sl], nvirt[sl]); + + ModuleBase::timer::end("OperatorLRHxc", "act"); } @@ -66,7 +73,7 @@ namespace LR this->pot.lock()->cal_v_eff(rho_trans, ucell, vr_hxc, ispin_ks); LR_Util::_deallocate_2order_nested_ptr(rho_trans, 1); - // 4. V^{Hxc}_{\mu,\nu}=\int{dr} \phi_\mu(r) v_{Hxc}(r) \phi_\mu(r) + // 4. V^{Hxc}_{\mu,\nu}=\int{dr} \phi_\mu(r) v_{Hxc}(r) \phi_\nu(r) this->hR->set_zero(); // clear hR for each bands ModuleGint::cal_gint_vl(vr_hxc.c, &*this->hR); ModuleBase::timer::end("OperatorLRHxc", "grid_calculation"); @@ -105,7 +112,7 @@ namespace LR LR_Util::_deallocate_2order_nested_ptr(rho_trans, 1); - // 4. V^{Hxc}_{\mu,\nu}=\int{dr} \phi_\mu(r) v_{Hxc}(r) \phi_\mu(r) + // 4. V^{Hxc}_{\mu,\nu}=\int{dr} \phi_\mu(r) v_{Hxc}(r) \phi_\nu(r) HR_real_imag.set_zero(); ModuleGint::cal_gint_vl(vr_hxc.c, &HR_real_imag); // LR_Util::print_HR(HR_real_imag, this->ucell.nat, "VR(real, 2d)"); diff --git a/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.h b/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.h index bb82780e14..2318056bfb 100644 --- a/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.h +++ b/source/source_lcao/module_lr/operator_casida/operator_lr_hxc.h @@ -33,6 +33,7 @@ namespace LR kv(kv_in), pX(pX_in), pc(pc_in), pmat(pmat_in), ispin_ks(ispin_ks) { ModuleBase::TITLE("OperatorLRHxc", "OperatorLRHxc"); + std::cout<<"Initializing OperatorLRHxc"<cal_type = hamilt::calculation_type::lcao_gint; this->is_first_node = true; this->hR = std::unique_ptr>(new hamilt::HContainer(&pmat_in)); @@ -58,7 +59,7 @@ namespace LR const int& nspin; const int& naos; const int nk = 1; - // const int nloc_per_band = 1; ///< local size of each state of X (passed by nbasis in act()) + // const int nloc_per_state = 1; ///< local size of each state of X (passed by nbasis in act()) const std::vector& nocc; const std::vector& nvirt; const std::vector ispin_ks = { 0 }; ///< the index of spin of psi_ks used in {AX, DM_trans} diff --git a/source/source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h b/source/source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h index f1bc110046..d2f497878e 100644 --- a/source/source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h +++ b/source/source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h @@ -15,13 +15,11 @@ namespace RI_Benchmark const int& nvirt, const psi::Psi& psi_ks_in, const TLRI& Cs_ao, - const TLRI& Vs, - const bool& read_from_aims=false, - const std::vector& aims_nbasis={}) + const TLRI& Vs) : naos(naos), nocc(nocc), nvirt(nvirt), npairs(nocc* nvirt), psi_ks(psi_ks_in), Cs_ao(Cs_ao), Vs(Vs), - Cs_vo_mo(cal_Cs_mo(ucell, Cs_ao, psi_ks_in, nocc, nvirt, /*occ_first=*/false, read_from_aims, aims_nbasis)), - Cs_ov_mo(cal_Cs_mo(ucell, Cs_ao, psi_ks_in, nocc, nvirt, /*occ_first=*/true, read_from_aims, aims_nbasis)), + Cs_vo_mo(cal_Cs_mo(ucell, Cs_ao, psi_ks_in, nocc, nvirt, /*occ_first=*/false)), + Cs_ov_mo(cal_Cs_mo(ucell, Cs_ao, psi_ks_in, nocc, nvirt, /*occ_first=*/true)), CV_vo(cal_CV(Cs_vo_mo, Vs)), CV_ov(cal_CV(Cs_ov_mo, Vs)) { @@ -39,7 +37,7 @@ namespace RI_Benchmark { Amat[i] = Amat1[i] + Amat2[i] + Amat3[i] + Amat4[i]; } - std::cout << "Amat_full (Hartree term) from RI (Unit Hartree):" << std::endl; + std::cout << "Amat_full (Hartree term) from RI (Unit Ry):" << std::endl;// cal_Amat_full has converted unit for (int i = 0;i < npairs;++i) { for (int j = 0;j < npairs;++j) diff --git a/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.h b/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.h index c4f5e741d8..0615bfe1df 100644 --- a/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.h +++ b/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.h @@ -2,10 +2,12 @@ #pragma once #include "source_cell/unitcell.h" #include "source_psi/psi.h" - +#include "source_lcao/module_lr/utils/lr_io.h" +#include "source_basis/module_ao/parallel_orbitals.h" #include namespace RI_Benchmark { + using TA = int; using TC = std::array; using TAC = std::pair; @@ -29,9 +31,7 @@ namespace RI_Benchmark const psi::Psi& wfc_ks, const int& nocc, const int& nvirt, - const int& occ_first=false, - const bool& read_from_aims=false, - const std::vector& aims_nbasis={}); + const int& occ_first=false); /// A=CVC, sum over atom quads template @@ -61,16 +61,13 @@ namespace RI_Benchmark TK* AX, const double& scale = 2.0); + // 3. read/write tools template - std::vector read_bands(const std::string& file, const int nocc, const int nvirt, int& ncore); + std::vector read_aims_ebands(const std::string& file, const int nocc, const int nvirt, int& ncore); + template void read_aims_eigenvectors(psi::Psi& wfc_ks, const std::string& file, const int ncore, const int nbands, const int nbasis); - /// only for blocking by atom pairs - template - TLRI read_coulomb_mat(const std::string& file, const TLRI& Cs); - /// for any way of blocking - template - TLRI read_coulomb_mat_general(const std::string& file, const TLRI& Cs); + template bool compare_Vs(const TLRI& Vs1, const TLRI& Vs2, const double thr = 1e-4); template diff --git a/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.hpp b/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.hpp index 501e3cf5f3..a7306047ed 100644 --- a/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.hpp +++ b/source/source_lcao/module_lr/ri_benchmark/ri_benchmark.hpp @@ -1,6 +1,10 @@ #pragma once +#include +#include #include "ri_benchmark.h" #include "source_base/module_container/base/third_party/blas.h" +#include "source_psi/psi.h" +#include "source_base/module_external/scalapack_connector.h" namespace RI_Benchmark { // std::cout << "the size of Cs:" << std::endl; @@ -51,28 +55,25 @@ namespace RI_Benchmark const psi::Psi& wfc_ks, const int& nocc, const int& nvirt, - const int& occ_first, - const bool& read_from_aims, - const std::vector& aims_nbasis) + const int& occ_first) { // assert(wfc_ks.get_nk() == 1); // currently only gamma-only is supported assert(nocc + nvirt <= wfc_ks.get_nbands()); - const bool use_aims_nbasis = (read_from_aims && !aims_nbasis.empty()); TLRI Cs_mo; int iw1 = 0; for (auto& c1 : Cs_ao) { const int& iat1 = c1.first; const int& it1 = ucell.iat2it[iat1]; - const int& nw1 = (use_aims_nbasis ? aims_nbasis[it1] : ucell.atoms[it1].nw); - if (!use_aims_nbasis) { assert(iw1 == ucell.get_iat2iwt()[iat1]); } + const int& nw1 = ucell.atoms[it1].nw; + assert(iw1 == ucell.get_iat2iwt()[iat1]); int iw2 = 0; for (auto& c2 : c1.second) { const int& iat2 = c2.first.first; const int& it2 = ucell.iat2it[iat2]; - const int& nw2 = (use_aims_nbasis ? aims_nbasis[it2] : ucell.atoms[it2].nw); - if (!use_aims_nbasis) { assert(iw2 == ucell.get_iat2iwt()[iat2]); } + const int& nw2 = ucell.atoms[it2].nw; + assert(iw2 == ucell.get_iat2iwt()[iat2]); const auto& tensor_ao = c2.second; const size_t& nabf = tensor_ao.shape[0]; @@ -118,7 +119,6 @@ namespace RI_Benchmark auto& Cs_shape = Cs_a.at(0).begin()->second.shape; auto& Vs_shape = Vs.at(0).begin()->second.shape; assert(Cs_shape.size() == 3); // abf, nocc, nvirt - assert(Cs_shape.size() == 3); // abf, nocc, nvirt assert(Vs_shape.size() == 2); // abf, abf const int& npairs = Cs_shape[1] * Cs_shape[2]; @@ -304,6 +304,8 @@ namespace RI_Benchmark std::cout << std::endl; return bands_final; } + + /// @brief read the eigenvectors from FHI-aims, only for gamma_only and spin degenerate template void read_aims_eigenvectors(psi::Psi& wfc_ks, const std::string& file, const int ncore, const int nbands, const int nbasis) { @@ -347,104 +349,7 @@ namespace RI_Benchmark std::cout << std::endl; } } - template < typename TR> // only for blocking by atom pairs - TLRI read_coulomb_mat(const std::string& file, const TLRI& Cs) - { //for gamma_only, V(q)=V(R=0) - std::ifstream ifs; - ifs.open(file); - size_t nks = 0, nabf = 0, istart = 0, jstart = 0, iend = 0, jend = 0; - std::string tmp; - ifs >> nks;// nkstot=1 - if (nks > 1) { std::cout << "Warning: nks>1 is not supported yet!" << std::endl; } - TLRI Vs; - const int nat = Cs.size(); - for (int iat1 = 0;iat1 < nat;++iat1) - { - const size_t nabf1 = Cs.at(iat1).at({ 0, {0,0,0} }).shape[0]; - for (int iat2 = 0;iat2 < nat;++iat2) - { - if (iat1 > iat2) - { // coulomb_mat has only the upper triangle part - Vs[iat1][{iat2, { 0,0,0 }}] = Vs[iat2][{iat1, { 0,0,0 }}].transpose(); - continue; - } - const size_t nabf2 = Cs.at(iat2).at({ 0, {0,0,0} }).shape[0]; - ifs >> nabf >> istart >> iend >> jstart >> jend >> tmp /*ik*/ >> tmp/*wk*/; - assert(nabf1 == iend - istart + 1); - assert(nabf2 == jend - jstart + 1); - RI::Tensor t({ nabf1, nabf2 }); - for (int i = 0;i < nabf1;++i) - { - for (int j = 0;j < nabf2;++j) - { - // t(i, j) = Vq[(istart + i) * nabf + jstart + j]; - ifs >> t(i, j) >> tmp; - } - } - Vs[iat1][{iat2, { 0,0,0 }}] = t; - } - } - return Vs; - } - - template < typename TR> // any blocking - TLRI read_coulomb_mat_general(const std::string& file, const TLRI& Cs) - { //for gamma_only, V(q)=V(R=0) - std::ifstream ifs; - ifs.open(file); - size_t nks = 0, nabf = 0, istart = 0, jstart = 0, iend = 0, jend = 0; - std::string tmp; - ifs >> nks;// nkstot=1 - if (nks > 1) { std::cout << "Warning: nks>1 is not supported yet!" << std::endl; } - TLRI Vs; - std::vector Vq; - while (ifs.peek() != EOF) - { - ifs >> nabf >> istart >> iend >> jstart >> jend >> tmp /*ik*/ >> tmp/*wk*/; - if (ifs.peek() == EOF) { break; } - if (Vq.empty()) { Vq.resize(nabf * nabf, 0.0); } - for (int i = istart - 1;i < iend;++i) - { - for (int j = jstart - 1;j < jend;++j) - { - ifs >> Vq[i * nabf + j] >> tmp; - } - } - } - const int nat = Cs.size(); - istart = 0; // - for (int iat1 = 0;iat1 < nat;++iat1) - { - const size_t nabf1 = Cs.at(iat1).at({ 0, {0,0,0} }).shape[0]; - jstart = 0; - for (int iat2 = 0;iat2 < nat;++iat2) - { - const size_t nabf2 = Cs.at(iat2).at({ 0, {0,0,0} }).shape[0]; - if (iat1 > iat2) - { // coulomb_mat has only the upper triangle part - Vs[iat1][{iat2, { 0,0,0 }}] = Vs[iat2][{iat1, { 0,0,0 }}].transpose(); - } - else - { - RI::Tensor t({ nabf1, nabf2 }); - for (int i = 0;i < nabf1;++i) - { - for (int j = 0;j < nabf2;++j) - { - t(i, j) = Vq[(istart + i) * nabf + jstart + j]; - } - } - Vs[iat1][{iat2, { 0,0,0 }}] = t; - } - jstart += nabf2; - } - assert(jstart == nabf); - istart += nabf1; - } - assert(istart == nabf); - return Vs; - } - + template < typename TR> bool compare_Vs(const TLRI& Vs1, const TLRI& Vs2, const double thr) { @@ -469,8 +374,9 @@ namespace RI_Benchmark } return true; } + // since ucell.nw is updated from aims_nbasis, this function is abandoned 25-05-23 template - std::vector> split_Ds(const std::vector>& Ds, const std::vector& aims_nbasis, const UnitCell& ucell) + std::vector> split_Ds(const std::vector>& Ds, const std::vector& aims_nbasis, const UnitCell& ucell) // vector index: ispin { // Due to the hard-coded constructor of elecstate::DensityMatrix, singlet-triplet with nspin=2 cannot use DM_trans with size 1 // if(Ds.size()>1) { throw std::runtime_error("split_Ds only supports gamma-only spin-1 Ds now."); } diff --git a/source/source_lcao/module_lr/utils/exciton_plotter.cpp b/source/source_lcao/module_lr/utils/exciton_plotter.cpp new file mode 100644 index 0000000000..c2be969f0d --- /dev/null +++ b/source/source_lcao/module_lr/utils/exciton_plotter.cpp @@ -0,0 +1,1042 @@ +#include "exciton_plotter.h" + +#include "source_base/constants.h" +#include "source_base/element_name.h" +#include "source_base/matrix3.h" +#include "source_base/module_external/blas_connector.h" +#include "source_base/parallel_global.h" +#include "source_base/ylm.h" +#include "source_cell/atom_spec.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +using Vec3 = ModuleBase::Vector3; + +Vec3 atom_position_bohr(const UnitCell& ucell, const int iat) +{ + return ucell.get_tau(iat) * ucell.lat0; +} + +int atomic_number_from_label(std::string element) +{ + element.erase(std::remove_if(element.begin(), element.end(), [](const char c) { return c >= '0' && c <= '9'; }), + element.end()); + for (int i = 0; i < static_cast(ModuleBase::element_name.size()); ++i) + { + if (element == ModuleBase::element_name[i]) + { + return i + 1; + } + } + return 0; +} + +std::array bvk_mesh_from_kpoints(const K_Vectors& kv, const int nk) +{ + if (kv.nmp[0] > 0 && kv.nmp[1] > 0 && kv.nmp[2] > 0) + { + return {kv.nmp[0], kv.nmp[1], kv.nmp[2]}; + } + + std::set kx, ky, kz; + const int n = std::min(nk, kv.kvec_d.size()); + constexpr double scale = 1.0e10; + for (int ik = 0; ik < n; ++ik) + { + kx.insert(static_cast(std::llround(kv.kvec_d[ik].x * scale))); + ky.insert(static_cast(std::llround(kv.kvec_d[ik].y * scale))); + kz.insert(static_cast(std::llround(kv.kvec_d[ik].z * scale))); + } + return {std::max(1, static_cast(kx.size())), + std::max(1, static_cast(ky.size())), + std::max(1, static_cast(kz.size()))}; +} + +struct SliceGeometry +{ + std::string plane; + double slice_pos = 0.0; + Vec3 a_bohr, b_bohr, c_bohr; + Vec3 u_vec, v_vec, perp_offset; + int u_axis = 0, v_axis = 1; + int nk_u = 1, nk_v = 1; + int u_start_cells = 0, u_end_cells = 1; + int v_start_cells = 0, v_end_cells = 1; + double u_start = 0.0, u_end = 1.0; + double v_start = 0.0, v_end = 1.0; + int res = 1, nu = 2, nv = 2; +}; + +SliceGeometry make_slice_geometry(const UnitCell& ucell, + const K_Vectors& kv, + const int nk, + const std::string& plane, + const double slice_pos, + const int npoints, + const double scale) +{ + SliceGeometry geom; + geom.plane = plane; + geom.slice_pos = slice_pos; + geom.a_bohr = ucell.a1 * ucell.lat0; + geom.b_bohr = ucell.a2 * ucell.lat0; + geom.c_bohr = ucell.a3 * ucell.lat0; + + const auto mesh = bvk_mesh_from_kpoints(kv, nk); + if (plane == "ab") + { + geom.u_vec = geom.a_bohr; + geom.v_vec = geom.b_bohr; + geom.perp_offset = geom.c_bohr * (slice_pos / geom.c_bohr.norm()); + geom.u_axis = 0; + geom.v_axis = 1; + geom.nk_u = mesh[0]; + geom.nk_v = mesh[1]; + } + else if (plane == "bc") + { + geom.u_vec = geom.b_bohr; + geom.v_vec = geom.c_bohr; + geom.perp_offset = geom.a_bohr * (slice_pos / geom.a_bohr.norm()); + geom.u_axis = 1; + geom.v_axis = 2; + geom.nk_u = mesh[1]; + geom.nk_v = mesh[2]; + } + else if (plane == "ca") + { + geom.u_vec = geom.c_bohr; + geom.v_vec = geom.a_bohr; + geom.perp_offset = geom.b_bohr * (slice_pos / geom.b_bohr.norm()); + geom.u_axis = 2; + geom.v_axis = 0; + geom.nk_u = mesh[2]; + geom.nk_v = mesh[0]; + } + else + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "Unknown slice plane: " + plane + ". Use ab, bc, or ca."); + } + + const double safe_scale = std::max(1.0, scale); + const int pad_u = std::max(0, static_cast(std::ceil((safe_scale - 1.0) * 0.5 * geom.nk_u))); + const int pad_v = std::max(0, static_cast(std::ceil((safe_scale - 1.0) * 0.5 * geom.nk_v))); + const int total_u = geom.nk_u + 2 * pad_u; + const int total_v = geom.nk_v + 2 * pad_v; + + geom.u_start_cells = static_cast(std::floor(0.5 - total_u * 0.5)); + geom.u_end_cells = geom.u_start_cells + total_u; + geom.v_start_cells = static_cast(std::floor(0.5 - total_v * 0.5)); + geom.v_end_cells = geom.v_start_cells + total_v; + geom.u_start = static_cast(geom.u_start_cells); + geom.u_end = static_cast(geom.u_end_cells); + geom.v_start = static_cast(geom.v_start_cells); + geom.v_end = static_cast(geom.v_end_cells); + + const int max_cells = std::max(total_u, total_v); + geom.res = std::max(1, npoints / std::max(1, max_cells)); + geom.nu = total_u * geom.res + 1; + geom.nv = total_v * geom.res + 1; + return geom; +} + +double bloch_phase_arg(const Vec3& kvec_d, const int n1, const int n2, const int n3) +{ + return ModuleBase::TWO_PI * (kvec_d.x * n1 + kvec_d.y * n2 + kvec_d.z * n3); +} + +double bloch_phase_arg_axis(const Vec3& kvec_d, const int axis, const int ncell) +{ + return ModuleBase::TWO_PI * kvec_d[axis] * ncell; +} + +std::complex to_complex(const double value) +{ + return std::complex(value, 0.0); +} + +std::complex to_complex(const std::complex& value) +{ + return value; +} + +char adjoint_flag(const double*) +{ + return 'T'; +} + +char adjoint_flag(const std::complex*) +{ + return 'C'; +} + +template +double density_from_dmk(const TK* dmk, const std::vector>& phi, const int naos) +{ + std::complex rho(0.0, 0.0); + for (int mu = 0; mu < naos; ++mu) + { + std::complex row_sum(0.0, 0.0); + for (int nu = 0; nu < naos; ++nu) + { + // cal_effective_dmk_* stores D^T for DensityMatrix; recover D(mu, nu) + // when contracting the matrix directly in real space. + row_sum += to_complex(dmk[nu + mu * naos]) * std::conj(phi[nu]); + } + rho += phi[mu] * row_sum; + } + return std::max(0.0, rho.real()); +} + +void write_slice_data(const UnitCell& ucell, + const SliceGeometry& geom, + const std::vector& density, + const std::string& filename, + const int istate, + const double energy_ry, + const std::string& density_kind, + const bool has_hole, + const std::array& r_h_fix) +{ + if (density.size() != static_cast(geom.nu * geom.nv)) + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "Invalid slice density dimensions."); + } + + std::ofstream ofs(filename); + if (!ofs) + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "Cannot open slice file: " + filename); + } + + ofs << std::setprecision(12); + ofs << "# ABACUS_EXCITON_SLICE 1\n"; + ofs << "# state " << istate << " energy_Ry " << energy_ry << " density_kind " << density_kind << "\n"; + const std::string fixed_particle + = density_kind == "conditional_hole" ? "electron" : (density_kind == "conditional_elec" ? "hole" : "none"); + ofs << "# fixed_particle " << fixed_particle; + if (has_hole) + { + ofs << " position_bohr " << r_h_fix[0] << " " << r_h_fix[1] << " " << r_h_fix[2]; + } + ofs << "\n"; + ofs << "# plane " << geom.plane << " slice_pos_bohr " << geom.slice_pos << "\n"; + ofs << "# grid " << geom.nu << " " << geom.nv << " u_range " << geom.u_start << " " << geom.u_end << " v_range " + << geom.v_start << " " << geom.v_end << "\n"; + ofs << "# bvk " << geom.nk_u << " " << geom.nk_v << "\n"; + ofs << "# lattice_bohr\n"; + ofs << "# a " << geom.a_bohr.x << " " << geom.a_bohr.y << " " << geom.a_bohr.z << "\n"; + ofs << "# b " << geom.b_bohr.x << " " << geom.b_bohr.y << " " << geom.b_bohr.z << "\n"; + ofs << "# c " << geom.c_bohr.x << " " << geom.c_bohr.y << " " << geom.c_bohr.z << "\n"; + ofs << "# axes_bohr\n"; + ofs << "# u " << geom.u_vec.x << " " << geom.u_vec.y << " " << geom.u_vec.z << "\n"; + ofs << "# v " << geom.v_vec.x << " " << geom.v_vec.y << " " << geom.v_vec.z << "\n"; + ofs << "# atoms " << ucell.nat << "\n"; + for (int iat = 0; iat < ucell.nat; ++iat) + { + const int it = ucell.iat2it[iat]; + const Vec3 tau = atom_position_bohr(ucell, iat); + ofs << "# " << ucell.atoms[it].label << " " << tau.x << " " << tau.y << " " << tau.z << "\n"; + } + ofs << "# data\n"; + ofs << std::scientific << std::setprecision(8); + for (int iu = 0; iu < geom.nu; ++iu) + { + for (int iv = 0; iv < geom.nv; ++iv) + { + ofs << density[iu * geom.nv + iv] << (iv + 1 < geom.nv ? " " : ""); + } + ofs << "\n"; + } +} +} // namespace + +// ============================================================================ +// OrbitalEvaluator non-template method implementations +// ============================================================================ + +void OrbitalEvaluator::init(const LCAO_Orbitals& orb, const UnitCell& ucell) +{ + atypes_.resize(ucell.ntype); + for (int it = 0; it < ucell.ntype; ++it) + { + const Numerical_Orbital& no = orb.Phi[it]; + const Atom& atom = ucell.atoms[it]; + auto& td = atypes_[it]; + + td.nw = atom.nw; + td.nwl = atom.nwl; + td.rcut = no.getRcut(); + td.dr_uniform = no.PhiLN(0, 0).dr_uniform; + td.radial_blocks.reserve(atom.nw); + + for (int iw = 0; iw < atom.nw; ++iw) + { + if (atom.iw2_new[iw]) + { + const int l = atom.iw2l[iw]; + const int n = atom.iw2n[iw]; + const auto& phi_ln = no.PhiLN(l, n); + + AtomTypeData::RadialBlock block; + block.begin_iw = iw; + block.size = 2 * l + 1; + block.ylm_begin = atom.iw2_ylm[iw]; + block.psi_uniform = phi_ln.psi_uniform.data(); + block.dpsi_uniform = phi_ln.dpsi_uniform.data(); + td.radial_blocks.push_back(block); + } + } + } +} + +/// @brief Replicates the cubic Hermite interpolation from GintAtom::set_phi() +void OrbitalEvaluator::eval_phi(const int it, const ModuleBase::Vector3& dr, double* phi_out) const +{ + const auto& td = atypes_[it]; + const double dist = dr.norm() < 1e-9 ? 1e-9 : dr.norm(); + if (dist > td.rcut) + { + ModuleBase::GlobalFunc::ZEROS(phi_out, td.nw); + return; + } + + std::vector ylma; + ModuleBase::Ylm::sph_harm(td.nwl, dr.x / dist, dr.y / dist, dr.z / dist, ylma); + + const double position = dist / td.dr_uniform; + const int ip = static_cast(position); + const double dx = position - ip; + const double dx2 = dx * dx; + const double dx3 = dx2 * dx; + const double c3 = 3.0 * dx2 - 2.0 * dx3; + const double c1 = 1.0 - c3; + const double c2 = (dx - 2.0 * dx2 + dx3) * td.dr_uniform; + const double c4 = (dx3 - dx2) * td.dr_uniform; + + const auto* blocks = td.radial_blocks.data(); + const int num_blocks = td.radial_blocks.size(); + for (int ib = 0; ib < num_blocks; ++ib) + { + const auto& block = blocks[ib]; + const double* psi_uniform = block.psi_uniform; + const double* dpsi_uniform = block.dpsi_uniform; + const double psi + = c1 * psi_uniform[ip] + c2 * dpsi_uniform[ip] + c3 * psi_uniform[ip + 1] + c4 * dpsi_uniform[ip + 1]; + + const int begin_iw = block.begin_iw; + const int end_iw = begin_iw + block.size; + int idx_lm = block.ylm_begin; + for (int iw = begin_iw; iw < end_iw; ++iw, ++idx_lm) + { + phi_out[iw] = psi * ylma[idx_lm]; + } + } +} + +// ============================================================================= +// Evaluate the Bloch-summed numerical atomic orbital at position r: +// phi^B_mu(r, k) = Sum_R exp(i * 2pi * kvec_d · (n1,n2,n3)) * phi_mu(r - tau - R) +// where kvec_d is the direct-coordinate k-vector. R runs over all lattice vectors +// within the orbital cutoff radius R = n1*a + n2*b + n3*c. +// ============================================================================= +void OrbitalEvaluator::eval_phi_all_bloch(const ModuleBase::Vector3& r, + const UnitCell& ucell, + std::complex* phi_bloch, + const ModuleBase::Vector3& kvec_d) const +{ + // Cell vectors in Bohr for real-space distance checks + const double lat0 = ucell.lat0; + ModuleBase::Vector3 a = ucell.a1 * lat0; + ModuleBase::Vector3 b = ucell.a2 * lat0; + ModuleBase::Vector3 c = ucell.a3 * lat0; + + // Determine how many lattice translations to check in each direction. + // nR = ceil(rcut / min(|a|,|b|,|c|)) + 1 ensures all overlapping orbitals are captured. + double max_rcut = 0.0; + for (const auto& td: atypes_) + max_rcut = std::max(max_rcut, td.rcut); + const int nR = static_cast(std::ceil(max_rcut / std::min({a.norm(), b.norm(), c.norm()}))) + 1; + + std::vector phi_tmp; + int iw_global = 0; + + // For each atom, iterate over neighboring cell images + for (int iat = 0; iat < ucell.nat; ++iat) + { + const int it = ucell.iat2it[iat]; + const int ia = ucell.iat2ia[iat]; + const ModuleBase::Vector3 tau = ucell.atoms[it].tau[ia] * lat0; + const int nw = atypes_[it].nw; + const double rcut = atypes_[it].rcut; + phi_tmp.resize(nw); + + // Triple loop over lattice vectors R = n1*a + n2*b + n3*c + for (int n1 = -nR; n1 <= nR; ++n1) + { + for (int n2 = -nR; n2 <= nR; ++n2) + { + for (int n3 = -nR; n3 <= nR; ++n3) + { + ModuleBase::Vector3 R = a * (double)n1 + b * (double)n2 + c * (double)n3; + ModuleBase::Vector3 dr = r - tau - R; + if (dr.norm() > rcut) + continue; // skip if orbital does not reach + + double kdotR = bloch_phase_arg(kvec_d, n1, n2, n3); + std::complex exp_ikR(std::cos(kdotR), std::sin(kdotR)); + + // Evaluate real-space orbital and accumulate with Bloch phase + eval_phi(it, dr, phi_tmp.data()); + for (int iw = 0; iw < nw; ++iw) + phi_bloch[iw_global + iw] += exp_ikR * phi_tmp[iw]; + } + } + } + iw_global += nw; + } +} + +namespace LR_Util +{ + +template +std::vector ExcitonPlotter::cal_effective_dmk_hole(const int istate) +{ + ModuleBase::TITLE("ExcitonPlotter", "cal_effective_dmk_hole"); + assert(this->nspin_x == 1); + const int offset_b = istate * this->ldim; + const int nocc = this->nocc[0]; + const int nvirt = this->nvirt[0]; + const auto data_type = container::DataTypeToEnum::value; + std::vector result(this->nk, + container::Tensor(data_type, DEV::CpuDevice, {this->naos, this->naos})); + const T alpha(1.0); + const T beta(0.0); + char adjoint = adjoint_flag(static_cast(nullptr)); + char normal = 'N'; + char transpose = 'T'; + + for (int ik = 0; ik < this->nk; ++ik) + { + this->psi_ks_vec[0].fix_k(ik); + T* const dmk = result[ik].data(); + ModuleBase::GlobalFunc::ZEROS(dmk, this->naos * this->naos); + const T* const xk = this->X + offset_b + ik * (this->ldim / this->nk); + + // Step 1: M_k = X_k^H * X_k → nocc x nocc Hermitian + container::Tensor M_k(data_type, DEV::CpuDevice, {nocc, nocc}); + T* const M_k_data = M_k.data(); + ModuleBase::GlobalFunc::ZEROS(M_k_data, nocc * nocc); + BlasConnector::gemm_cm(adjoint, normal, nocc, nocc, nvirt, + alpha, xk, nvirt, xk, nvirt, + beta, M_k_data, nocc); + + // Step 2: temp = M_k * C_occ^H → nocc x naos + container::Tensor temp(data_type, DEV::CpuDevice, {nocc, this->naos}); + T* const temp_data = temp.data(); + BlasConnector::gemm_cm(normal, adjoint, nocc, this->naos, nocc, + alpha, M_k_data, nocc, + this->psi_ks_vec[0].get_pointer(0), this->naos, + beta, temp_data, nocc); + + // Step 3: dmk = temp^T * C_occ^T = (C_occ * M_k * C_occ^H)^T. + // DensityMatrix expects the AO indices in this transposed order. + BlasConnector::gemm_cm(transpose, transpose, this->naos, this->naos, nocc, + alpha, temp_data, nocc, + this->psi_ks_vec[0].get_pointer(0), this->naos, + beta, dmk, this->naos); + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "cal effective dmk hole"); + return result; +} + +template +std::vector ExcitonPlotter::cal_effective_dmk_elec(const int istate) +{ + ModuleBase::TITLE("ExcitonPlotter", "cal_effective_dmk_elec"); + assert(this->nspin_x == 1); + const int offset_b = istate * this->ldim; + const int nocc = this->nocc[0]; + const int nvirt = this->nvirt[0]; + const auto data_type = container::DataTypeToEnum::value; + std::vector result(this->nk, + container::Tensor(data_type, DEV::CpuDevice, {this->naos, this->naos})); + const T alpha(1.0); + const T beta(0.0); + char adjoint = adjoint_flag(static_cast(nullptr)); + char normal = 'N'; + char transpose = 'T'; + + for (int ik = 0; ik < this->nk; ++ik) + { + this->psi_ks_vec[0].fix_k(ik); + T* const dmk = result[ik].data(); + ModuleBase::GlobalFunc::ZEROS(dmk, this->naos * this->naos); + const T* const xk = this->X + offset_b + ik * (this->ldim / this->nk); + + // Step 1: N_k = X_k * X_k^H → nvirt x nvirt Hermitian + container::Tensor N_k(data_type, DEV::CpuDevice, {nvirt, nvirt}); + T* const N_k_data = N_k.data(); + ModuleBase::GlobalFunc::ZEROS(N_k_data, nvirt * nvirt); + BlasConnector::gemm_cm(normal, adjoint, nvirt, nvirt, nocc, + alpha, xk, nvirt, xk, nvirt, + beta, N_k_data, nvirt); + + // Step 2: temp = N_k * C_virt^H → nvirt x naos + container::Tensor temp(data_type, DEV::CpuDevice, {nvirt, this->naos}); + T* const temp_data = temp.data(); + BlasConnector::gemm_cm(normal, adjoint, nvirt, this->naos, nvirt, + alpha, N_k_data, nvirt, + this->psi_ks_vec[0].get_pointer(nocc), this->naos, + beta, temp_data, nvirt); + + // Step 3: dmk = temp^T * C_virt^T = (C_virt * N_k * C_virt^H)^T. + // DensityMatrix expects the AO indices in this transposed order. + BlasConnector::gemm_cm(transpose, transpose, this->naos, this->naos, nvirt, + alpha, temp_data, nvirt, + this->psi_ks_vec[0].get_pointer(nocc), this->naos, + beta, dmk, this->naos); + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "cal effective dmk elec"); + return result; +} + +template +void ExcitonPlotter::plot_average_density(const int istate, const std::string& type) +{ + ModuleBase::TITLE("ExcitonPlotter", "plot_average_density"); + if (type != "hole" && type != "elec") + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "Unknown average density type: " + type + ". Use hole or elec."); + } + const auto dmk = type == "hole" ? cal_effective_dmk_hole(istate) : cal_effective_dmk_elec(istate); + elecstate::DensityMatrix dm(&this->pmat, this->nspin_x, this->kv.kvec_d, this->nk); + for (int ik = 0; ik < this->nk; ++ik) + { + dm.set_DMK_pointer(ik, dmk[ik].template data()); + } + LR_Util::initialize_DMR(dm, this->pmat, this->ucell, this->gd_, this->orb_cutoff_); + dm.cal_DMR(); + + double** rho_result = nullptr; + LR_Util::_allocate_2order_nested_ptr(rho_result, this->nspin_x, this->rho_basis.nrxx); + for (int is = 0; is < this->nspin_x; ++is) + { + ModuleBase::GlobalFunc::ZEROS(rho_result[is], this->rho_basis.nrxx); + } + ModuleGint::cal_gint_rho(dm.get_DMR_vector(), this->nspin_x, rho_result, false); + + for (int is = 0; is < this->nspin_x; ++is) + { + const std::string filename = this->output_dir_ + "Exciton_avg_" + type + "_state" + std::to_string(istate) + + "_spin" + std::to_string(is) + ".cube"; + ModuleIO::write_vdata_palgrid(this->Pgrid, + rho_result[is], + is, + this->nspin_x, + 0, + filename, + this->eig[istate], + &this->ucell, + 8, + 0, + false, /*two_fermi*/ + false); + } + LR_Util::_deallocate_2order_nested_ptr(rho_result, this->nspin_x); +} + +template +void ExcitonPlotter::plot_average_slice(const int istate, + const std::string& type, + const std::string& plane, + const double slice_pos, + const int npoints, + const double scale) +{ + ModuleBase::TITLE("ExcitonPlotter", "plot_average_slice"); + assert(this->nspin_x == 1); + if (!this->orb_) + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", + "plot_average_slice requires LCAO_Orbitals; pass orb to constructor."); + } + if (type != "hole" && type != "elec") + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "Unknown average density type: " + type + ". Use hole or elec."); + } + + const auto dmk = type == "hole" ? cal_effective_dmk_hole(istate) : cal_effective_dmk_elec(istate); + const SliceGeometry geom = make_slice_geometry(this->ucell, this->kv, this->nk, plane, slice_pos, npoints, scale); + const double du = (geom.u_end - geom.u_start) / (geom.nu - 1); + const double dv = (geom.v_end - geom.v_start) / (geom.nv - 1); + const int cell_res = geom.res; + + std::vector density(geom.nu * geom.nv, 0.0); + std::vector rho_cache(cell_res * cell_res, 0.0); + std::vector cache_valid(cell_res * cell_res, false); + std::vector> phi_bloch(this->naos); + + for (int iu = 0; iu < geom.nu; ++iu) + { + const double u_frac = geom.u_start + du * iu; + const double u_cell = u_frac - std::floor(u_frac); + int uc_idx = static_cast(std::round(u_cell * cell_res)); + if (uc_idx >= cell_res) + { + uc_idx -= cell_res; + } + + for (int iv = 0; iv < geom.nv; ++iv) + { + const double v_frac = geom.v_start + dv * iv; + const double v_cell = v_frac - std::floor(v_frac); + int vc_idx = static_cast(std::round(v_cell * cell_res)); + if (vc_idx >= cell_res) + { + vc_idx -= cell_res; + } + + const int cache_idx = uc_idx * cell_res + vc_idx; + if (!cache_valid[cache_idx]) + { + const Vec3 position = geom.perp_offset + geom.u_vec * u_cell + geom.v_vec * v_cell; + double rho = 0.0; + for (int ik = 0; ik < this->nk; ++ik) + { + std::fill(phi_bloch.begin(), phi_bloch.end(), std::complex(0.0, 0.0)); + const Vec3 kvec_d + = ik < static_cast(this->kv.kvec_d.size()) ? this->kv.kvec_d[ik] : Vec3(0.0, 0.0, 0.0); + this->orb_eval_.eval_phi_all_bloch(position, this->ucell, phi_bloch.data(), kvec_d); + rho += density_from_dmk(dmk[ik].template data(), phi_bloch, this->naos); + } + rho_cache[cache_idx] = rho; + cache_valid[cache_idx] = true; + } + density[iu * geom.nv + iv] = rho_cache[cache_idx]; + } + } + + const std::string filename + = this->output_dir_ + "Exciton_avg_" + type + "_slice_state" + std::to_string(istate) + ".dat"; + write_slice_data(this->ucell, + geom, + density, + filename, + istate, + this->eig[istate], + "average_" + type, + false, + {0.0, 0.0, 0.0}); + std::cout << "Average " << type << " density slice written to " << filename << " (" << geom.nu << "x" << geom.nv + << ", " << geom.res << " pts/cell)" << std::endl; +} + +template +std::vector>> ExcitonPlotter::build_conditional_coefficients( + const int istate, + const std::array& r_fix_in, + const bool plot_electron) +{ + using Complex = std::complex; + const Vec3 r_fix(r_fix_in[0], r_fix_in[1], r_fix_in[2]); + const int nocc = this->nocc[0]; + const int nvirt = this->nvirt[0]; + const int nmix = plot_electron ? nvirt : nocc; + const int offset_b = istate * this->ldim; + + std::vector> mixing(this->nk, std::vector(nmix, Complex(0.0, 0.0))); + for (int ik = 0; ik < this->nk; ++ik) + { + const int x_start = ik * nocc * nvirt; + if (plot_electron) + { + for (int v = 0; v < nocc; ++v) + { + const Complex fixed_wfc = std::conj( + this->orb_eval_ + .eval_wfc_bloch(r_fix, ik, v, this->psi_ks_vec[0], this->ucell, this->kv.kvec_d[ik])); + for (int c = 0; c < nvirt; ++c) + { + mixing[ik][c] += to_complex(this->X[offset_b + x_start + c + v * nvirt]) * fixed_wfc; + } + } + } + else + { + for (int c = 0; c < nvirt; ++c) + { + const Complex fixed_wfc = this->orb_eval_.eval_wfc_bloch(r_fix, + ik, + nocc + c, + this->psi_ks_vec[0], + this->ucell, + this->kv.kvec_d[ik]); + for (int v = 0; v < nocc; ++v) + { + mixing[ik][v] += to_complex(this->X[offset_b + x_start + c + v * nvirt]) * fixed_wfc; + } + } + } + } + + std::vector> coefficients(this->nk, std::vector(this->naos, Complex(0.0, 0.0))); + for (int ik = 0; ik < this->nk; ++ik) + { + this->psi_ks_vec[0].fix_k(ik); + if (!plot_electron) + { + for (Complex& value: mixing[ik]) + { + value = std::conj(value); + } + } + const int band_start = plot_electron ? nocc : 0; + for (int ib = 0; ib < nmix; ++ib) + { + const T* const band = this->psi_ks_vec[0].get_pointer(band_start + ib); + for (int mu = 0; mu < this->naos; ++mu) + { + coefficients[ik][mu] += to_complex(band[mu]) * mixing[ik][ib]; + } + } + } + return coefficients; +} + +template +void ExcitonPlotter::plot_conditional_density(const int istate, + const std::array& r_fix_in, + const std::string& type) +{ + using Complex = std::complex; + ModuleBase::TITLE("ExcitonPlotter", "plot_conditional_density"); + assert(this->nspin_x == 1); + if (!this->orb_) + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", + "plot_conditional_density requires LCAO_Orbitals; pass orb to constructor."); + } + const bool plot_electron = type == "elec"; + if (!plot_electron && type != "hole") + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "Unknown conditional density type: " + type + ". Use elec or hole."); + } + const auto coefficients = build_conditional_coefficients(istate, r_fix_in, plot_electron); + + const int nx = this->rho_basis.nx; + const int ny = this->rho_basis.ny; + const int nz = this->rho_basis.nz; + const int nrxx = this->rho_basis.nrxx; + const Vec3 a = this->ucell.a1 * this->ucell.lat0; + const Vec3 b = this->ucell.a2 * this->ucell.lat0; + const Vec3 c = this->ucell.a3 * this->ucell.lat0; + const Vec3 dx = a * (1.0 / nx); + const Vec3 dy = b * (1.0 / ny); + const Vec3 dz = c * (1.0 / nz); + + const auto mesh = bvk_mesh_from_kpoints(this->kv, this->nk); + const int nk1 = mesh[0]; + const int nk2 = mesh[1]; + const int nk3 = mesh[2]; + const bool write_supercell = nk1 > 1 || nk2 > 1 || nk3 > 1; + const int snx = nx * nk1; + const int sny = ny * nk2; + const int snz = nz * nk3; + + double** rho_result = nullptr; + LR_Util::_allocate_2order_nested_ptr(rho_result, this->nspin_x, nrxx); + ModuleBase::GlobalFunc::ZEROS(rho_result[0], nrxx); + std::cout << "Conditional density: coherent sum over " << this->nk << " k-points" + << ", BvK supercell " << nk1 << "x" << nk2 << "x" << nk3 << ", grid " << snx << "x" << sny << "x" << snz + << std::endl; + + std::vector supercell_data; + std::vector> psi_k_cache; + if (write_supercell) + { + supercell_data.resize(snx * sny * snz, 0.0); + psi_k_cache.resize(this->nk, std::vector(nrxx)); + } + +#ifdef _OPENMP +#pragma omp parallel for schedule(dynamic) +#endif + for (int ix = 0; ix < nx; ++ix) + { + std::vector phi_bloch(this->naos); + for (int iy = 0; iy < ny; ++iy) + { + for (int iz = 0; iz < nz; ++iz) + { + const int index = ix * ny * nz + iy * nz + iz; + const Vec3 position + = dx * static_cast(ix) + dy * static_cast(iy) + dz * static_cast(iz); + Complex conditional_wfc(0.0, 0.0); + for (int ik = 0; ik < this->nk; ++ik) + { + std::fill(phi_bloch.begin(), phi_bloch.end(), Complex(0.0, 0.0)); + this->orb_eval_.eval_phi_all_bloch(position, this->ucell, phi_bloch.data(), this->kv.kvec_d[ik]); + Complex wfc_k(0.0, 0.0); + for (int mu = 0; mu < this->naos; ++mu) + { + wfc_k += coefficients[ik][mu] * phi_bloch[mu]; + } + if (!plot_electron) + { + wfc_k = std::conj(wfc_k); + } + if (write_supercell) + { + psi_k_cache[ik][index] = wfc_k; + } + conditional_wfc += wfc_k; + } + rho_result[0][index] = std::norm(conditional_wfc); + } + } + } + std::cout << "Home cell coherent evaluation complete." << std::endl; + + const std::string home_filename + = this->output_dir_ + "Exciton_cond_" + type + "_state" + std::to_string(istate) + "_spin0_home.cube"; + ModuleIO::write_vdata_palgrid(this->Pgrid, + rho_result[0], + 0, + this->nspin_x, + 0, + home_filename, + this->eig[istate], + &this->ucell, + 8, + 0, + false, /*two_fermi*/ + false); + + if (write_supercell) + { + std::cout << "Writing BvK supercell cube (" << snx << "x" << sny << "x" << snz << ")..." << std::endl; +#ifdef _OPENMP +#pragma omp parallel for collapse(3) schedule(dynamic) +#endif + for (int ni = 0; ni < nk1; ++ni) + { + for (int nj = 0; nj < nk2; ++nj) + { + for (int nk_cell = 0; nk_cell < nk3; ++nk_cell) + { + for (int ix = 0; ix < nx; ++ix) + { + for (int iy = 0; iy < ny; ++iy) + { + for (int iz = 0; iz < nz; ++iz) + { + const int home_index = ix * ny * nz + iy * nz + iz; + Complex wfc(0.0, 0.0); + for (int ik = 0; ik < this->nk; ++ik) + { + double phase_arg = bloch_phase_arg(this->kv.kvec_d[ik], ni, nj, nk_cell); + if (!plot_electron) + { + phase_arg = -phase_arg; + } + const Complex phase(std::cos(phase_arg), std::sin(phase_arg)); + wfc += phase * psi_k_cache[ik][home_index]; + } + const int super_index + = (ni * nx + ix) * sny * snz + (nj * ny + iy) * snz + nk_cell * nz + iz; + supercell_data[super_index] = std::norm(wfc); + } + } + } + } + } + } + + const Vec3 super_a = a * static_cast(nk1); + const Vec3 super_b = b * static_cast(nk2); + const Vec3 super_c = c * static_cast(nk3); + const int super_natoms = this->ucell.nat * nk1 * nk2 * nk3; + std::vector atom_type; + std::vector atom_charge; + std::vector> atom_position; + atom_type.reserve(super_natoms); + atom_charge.reserve(super_natoms); + atom_position.reserve(super_natoms); + for (int iat = 0; iat < this->ucell.nat; ++iat) + { + const int it = this->ucell.iat2it[iat]; + const Vec3 tau = atom_position_bohr(this->ucell, iat); + for (int ni = 0; ni < nk1; ++ni) + { + for (int nj = 0; nj < nk2; ++nj) + { + for (int nk_cell = 0; nk_cell < nk3; ++nk_cell) + { + const Vec3 position = tau + a * static_cast(ni) + b * static_cast(nj) + + c * static_cast(nk_cell); + atom_type.push_back(atomic_number_from_label(this->ucell.atoms[it].label)); + atom_charge.push_back(this->ucell.atoms[it].ncpp.zv); + atom_position.push_back({position.x, position.y, position.z}); + } + } + } + } + + const std::string supercell_filename + = this->output_dir_ + "Exciton_cond_" + type + "_state" + std::to_string(istate) + "_spin0_supercell.cube"; + ModuleIO::write_cube(supercell_filename, + {"BvK supercell conditional " + type + " density, state " + std::to_string(istate), + "Exciton energy (Ry) = " + std::to_string(this->eig[istate])}, + super_natoms, + {0.0, 0.0, 0.0}, + snx, + sny, + snz, + {super_a.x / snx, super_a.y / snx, super_a.z / snx}, + {super_b.x / sny, super_b.y / sny, super_b.z / sny}, + {super_c.x / snz, super_c.y / snz, super_c.z / snz}, + atom_type, + atom_charge, + atom_position, + supercell_data, + 8); + std::cout << "Supercell cube written to " << supercell_filename << std::endl; + } + + LR_Util::_deallocate_2order_nested_ptr(rho_result, this->nspin_x); + std::cout << "Conditional " << type << " density for state " << istate << ", " + << (plot_electron ? "hole" : "electron") << " fixed at (" << r_fix_in[0] << ", " << r_fix_in[1] << ", " + << r_fix_in[2] << ") Bohr" << std::endl; +} + +template +void ExcitonPlotter::plot_cond_slice(const int istate, + const std::array& r_fix_in, + const std::string& plane, + const double slice_pos, + const int npoints, + const double scale, + const std::string& type) +{ + using Complex = std::complex; + ModuleBase::TITLE("ExcitonPlotter", "plot_cond_slice"); + assert(this->nspin_x == 1); + if (!this->orb_) + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "plot_cond_slice requires LCAO_Orbitals; pass orb to constructor."); + } + const bool plot_electron = type == "elec"; + if (!plot_electron && type != "hole") + { + ModuleBase::WARNING_QUIT("ExcitonPlotter", "Unknown conditional slice type: " + type + ". Use elec or hole."); + } + const auto coefficients = build_conditional_coefficients(istate, r_fix_in, plot_electron); + const SliceGeometry geom = make_slice_geometry(this->ucell, this->kv, this->nk, plane, slice_pos, npoints, scale); + const double du = (geom.u_end - geom.u_start) / (geom.nu - 1); + const double dv = (geom.v_end - geom.v_start) / (geom.nv - 1); + + std::vector phase_u(this->nk); + std::vector phase_v(this->nk); + for (int ik = 0; ik < this->nk; ++ik) + { + phase_u[ik] = bloch_phase_arg_axis(this->kv.kvec_d[ik], geom.u_axis, 1); + phase_v[ik] = bloch_phase_arg_axis(this->kv.kvec_d[ik], geom.v_axis, 1); + } + + const int cell_res = geom.res; + std::vector density(geom.nu * geom.nv, 0.0); + std::vector phi_bloch(this->naos); + std::vector> psi_cache(cell_res * cell_res, std::vector(this->nk, Complex(0.0, 0.0))); + std::vector cache_valid(cell_res * cell_res, false); + + for (int iu = 0; iu < geom.nu; ++iu) + { + const double u_frac = geom.u_start + du * iu; + const double u_cell = u_frac - std::floor(u_frac); + const int u_translation = static_cast(std::floor(u_frac)); + int uc_idx = static_cast(std::round(u_cell * cell_res)); + if (uc_idx >= cell_res) + { + uc_idx -= cell_res; + } + + for (int iv = 0; iv < geom.nv; ++iv) + { + const double v_frac = geom.v_start + dv * iv; + const double v_cell = v_frac - std::floor(v_frac); + const int v_translation = static_cast(std::floor(v_frac)); + int vc_idx = static_cast(std::round(v_cell * cell_res)); + if (vc_idx >= cell_res) + { + vc_idx -= cell_res; + } + + const int cache_index = uc_idx * cell_res + vc_idx; + if (!cache_valid[cache_index]) + { + const Vec3 home_position = geom.perp_offset + geom.u_vec * u_cell + geom.v_vec * v_cell; + for (int ik = 0; ik < this->nk; ++ik) + { + std::fill(phi_bloch.begin(), phi_bloch.end(), Complex(0.0, 0.0)); + this->orb_eval_.eval_phi_all_bloch(home_position, + this->ucell, + phi_bloch.data(), + this->kv.kvec_d[ik]); + Complex wfc_k(0.0, 0.0); + for (int mu = 0; mu < this->naos; ++mu) + { + wfc_k += coefficients[ik][mu] * phi_bloch[mu]; + } + psi_cache[cache_index][ik] = plot_electron ? wfc_k : std::conj(wfc_k); + } + cache_valid[cache_index] = true; + } + + Complex conditional_wfc(0.0, 0.0); + for (int ik = 0; ik < this->nk; ++ik) + { + double phase_arg = phase_u[ik] * u_translation + phase_v[ik] * v_translation; + if (!plot_electron) + { + phase_arg = -phase_arg; + } + const Complex phase(std::cos(phase_arg), std::sin(phase_arg)); + conditional_wfc += phase * psi_cache[cache_index][ik]; + } + density[iu * geom.nv + iv] = std::norm(conditional_wfc); + } + } + + std::cout << "Bloch-sum evaluation complete (cached " << cell_res << "x" << cell_res << " home-cell positions)." + << std::endl; + const std::string filename + = this->output_dir_ + "Exciton_cond_" + type + "_slice_state" + std::to_string(istate) + ".dat"; + write_slice_data(this->ucell, + geom, + density, + filename, + istate, + this->eig[istate], + "conditional_" + type, + true, + r_fix_in); + std::cout << "Conditional density slice written to " << filename << " (" << geom.nu << "x" << geom.nv << ", range " + << scale << "x BvK, " << geom.res << " pts/cell)" << std::endl; +} + +template class ExcitonPlotter; +template class ExcitonPlotter>; + +} // namespace LR_Util diff --git a/source/source_lcao/module_lr/utils/exciton_plotter.h b/source/source_lcao/module_lr/utils/exciton_plotter.h new file mode 100644 index 0000000000..f49126820c --- /dev/null +++ b/source/source_lcao/module_lr/utils/exciton_plotter.h @@ -0,0 +1,265 @@ +#pragma once +#include "source_base/tool_title.h" +#include "source_base/ylm.h" +#include "source_basis/module_ao/ORB_read.h" +#include "source_cell/atom_spec.h" +#include "source_cell/klist.h" +#include "source_estate/module_dm/density_matrix.h" +#include "source_io/module_output/cube_io.h" +#include "source_lcao/module_gint/gint_interface.h" +#include "source_lcao/module_lr/dm_trans/dm_trans.h" +#include "source_lcao/module_lr/utils/lr_util.h" +#include "source_lcao/module_lr/utils/lr_util_hcontainer.h" + +#include +#include +#include +#include + +// ============================================================================ +// OrbitalEvaluator: single-point NAO wavefunction evaluation +// Replicates the cubic Hermite interpolation from GintAtom::set_phi() +// ============================================================================ +struct OrbitalEvaluator +{ + struct AtomTypeData + { + struct RadialBlock + { + int begin_iw = 0; + int size = 0; + int ylm_begin = 0; + const double* psi_uniform = nullptr; + const double* dpsi_uniform = nullptr; + }; + + int nw = 0, nwl = 0; ///< number of NAOs, max L for this atom type + double rcut = 0.0, dr_uniform = 0.0; ///< orbital cutoff radius, uniform grid spacing + std::vector radial_blocks; + }; + std::vector atypes_; + + /// @brief Initialize evaluator by extracting radial function pointers from LCAO_Orbitals + /// @param orb numerical atomic orbital data (one Numerical_Orbital per atom type) + void init(const LCAO_Orbitals& orb, const UnitCell& ucell); + + /// @brief Evaluate all NAO values for one atom type at relative position dr (in Bohr) + /// @param it atom type index + /// @param dr Cartesian vector from atom center to evaluation point (Bohr) + /// @param phi_out array of length nw[it], filled with phi values + void eval_phi(int it, const ModuleBase::Vector3& dr, double* phi_out) const; + + /// @brief Evaluate Bloch-summed NAO values including neighboring cell contributions + /// phi^B_mu(r, k) = Sum_R exp(i * 2pi * k_d·R_int) * phi_mu(r - tau - R). + /// The R-sum runs over all lattice vectors within the orbital cutoff radius. + /// Output is complex because the Bloch phase contributes a complex factor. + /// The caller must zero-initialize phi_bloch before calling (this function accumulates). + /// @param r evaluation point in Bohr + /// @param ucell unit cell + /// @param phi_bloch output array of length total_naos (complex), values are added + /// @param kvec_d direct-coordinate k-vector + void eval_phi_all_bloch(const ModuleBase::Vector3& r, + const UnitCell& ucell, + std::complex* phi_bloch, + const ModuleBase::Vector3& kvec_d) const; + + /// @brief Evaluate a Bloch-summed LCAO wavefunction at one point. + /// Uses the same neighboring-cell AO sum as eval_phi_all_bloch(). + template + std::complex eval_wfc_bloch(const ModuleBase::Vector3& r_fix, + int ik, + int iband, + const psi::Psi& psi_ks, + const UnitCell& ucell, + const ModuleBase::Vector3& kvec_d) const + { + psi_ks.fix_k(ik); + int total_naos = 0; + for (int iat = 0; iat < ucell.nat; ++iat) + { + total_naos += atypes_[ucell.iat2it[iat]].nw; + } + + std::vector> phi_bloch(total_naos, std::complex(0.0, 0.0)); + eval_phi_all_bloch(r_fix, ucell, phi_bloch.data(), kvec_d); + + std::complex result(0.0, 0.0); + for (int iw = 0; iw < total_naos; ++iw) + { + result += std::complex(psi_ks(iband, iw)) * phi_bloch[iw]; + } + return result; + } +}; + +namespace LR_Util +{ +template +class ExcitonPlotter +{ + /// @brief Construct an exciton density plotter for TDA excitons. + /// Supports average density (integrating out one particle coordinate) and + /// conditional density (fixing the hole position and evaluating the coherent + /// electron wavefunction). + /// @param nspin_global global number of spin channels + /// @param naos number of atomic-orbital basis functions + /// @param nocc number of occupied bands per spin + /// @param nvirt number of virtual bands per spin + /// @param psi_ks_in Kohn-Sham wavefunction coefficients + /// @param ucell_in unit cell + /// @param kv_in k-point list + /// @param gd_in grid driver for neighbor-list search + /// @param orb_cutoff_in orbital cutoff radii per atom type + /// @param Pgrid_in parallel grid for real-space decomposition + /// @param rho_basis_in plane-wave basis for the real-space grid + /// @param pX_in 2D block-cyclic distribution for X matrix + /// @param pc_in 2D block-cyclic distribution for KS orbitals (reserved for distributed plotting) + /// @param pmat_in parallel distribution for AO matrices + /// @param output_dir_in output directory for generated density files + /// @param two_fermi_in whether spin channels use separate Fermi energies + /// @param eig excitation energies in Ry + /// @param X BSE eigenvector amplitudes (TDA) + /// @param openshell true for spin-unrestricted (nspin_x=2) + /// @param orb optional LCAO_Orbitals pointer (required for conditional density) + public: + ExcitonPlotter(const int nspin_global, + const int naos, + const std::vector& nocc, + const std::vector& nvirt, + const psi::Psi& psi_ks_in, + const UnitCell& ucell_in, + const K_Vectors& kv_in, + const Grid_Driver& gd_in, + const std::vector& orb_cutoff_in, + const Parallel_Grid& Pgrid_in, + const ModulePW::PW_Basis& rho_basis_in, + const std::vector& pX_in, + const Parallel_2D& pc_in, + const Parallel_Orbitals& pmat_in, + const std::string& output_dir_in, + const double* eig, + const T* X, + const bool openshell, + const LCAO_Orbitals* orb) + : nspin_x(openshell ? 2 : 1), naos(naos), nocc(nocc), nvirt(nvirt), + nk(nspin_global == 2 ? kv_in.get_nks() / 2 : kv_in.get_nks()), + ldim(nk * (nspin_x == 2 ? pX_in[0].get_local_size() + pX_in[1].get_local_size() : pX_in[0].get_local_size())), + eig(eig), X(X), kv(kv_in), ucell(ucell_in), orb_cutoff_(orb_cutoff_in), gd_(gd_in), Pgrid(Pgrid_in), + rho_basis(rho_basis_in), output_dir_(output_dir_in), pc(pc_in), pmat(pmat_in), orb_(orb) + { + for (int is = 0; is < nspin_global; ++is) + { + psi_ks_vec.emplace_back(LR_Util::get_psi_spin(psi_ks_in, is, nk)); + } + if (orb) + { + orb_eval_.init(*orb, ucell_in); + } + }; + // --------------------------------------------------------------------- + // The average density integrates out one particle coordinate. + // The conditional density fixes the hole position and evaluates the coherent electron wavefunction. + // + // Both use Hermitian effective DMK → DMR → cal_gint_rho (average) or + // direct coherent real-space summation (conditional) to produce proper + // non-negative densities. + + /// @brief Compute effective DMK for average hole density + /// Produces D_hole(k)^T, where D_hole(k) = C_occ(k) * (X_k^H * X_k) * C_occ(k)^H, + /// matching the AO-index order expected by DensityMatrix::cal_DMR(). + /// Marginalizes over conduction bands: M_k[v,v'] = Sum_c A_{kvc} * conj(A_{kv'c}). + /// @param istate BSE state index + /// @return dmk_per_kpoint + std::vector cal_effective_dmk_hole(const int istate); + + /// @brief Compute effective DMK for average electron density + /// Produces D_elec(k)^T, where D_elec(k) = C_virt(k) * (X_k * X_k^H) * C_virt(k)^H, + /// matching the AO-index order expected by DensityMatrix::cal_DMR(). + /// Marginalizes over valence bands: N_k[c,c'] = Sum_v A_{kvc} * conj(A_{kvc'}). + /// @param istate BSE state index + /// @return dmk_per_kpoint + std::vector cal_effective_dmk_elec(const int istate); + + /// @brief Plot average hole or electron density as a .cube file + /// Integrates out the other particle coordinate. The effective DMK is Hermitian, + /// so DMR's real part suffices for cal_gint_rho — no squaring needed. + /// @param istate BSE state index + /// @param type "hole" for average hole density, "elec" for average electron density + void plot_average_density(const int istate, const std::string& type); + + /// @brief Plot average hole or electron density on a 2D cross-section + /// Evaluates the Hermitian effective density matrix directly on the requested + /// plane and writes a data file readable by plot_cond_silce.py. + /// @param istate BSE state index + /// @param type "hole" for average hole density, "elec" for average electron density + /// @param plane cross-section plane: "ab", "bc", or "ca" + /// @param slice_pos offset along the perpendicular direction (Bohr) + /// @param npoints desired grid resolution + /// @param scale view range relative to the plotted periodic cell + void plot_average_slice(const int istate, + const std::string& type, + const std::string& plane, + double slice_pos, + int npoints, + double scale); + + /// @brief Plot conditional electron or hole density as home-cell and BvK .cube files + /// Fixes the opposite particle at r_fix and evaluates the coherent conditional + /// wavefunction. Uses direct real-space evaluation (NOT the DMR pipeline) to + /// preserve k-point coherence. + /// @param istate BSE state index + /// @param r_fix fixed opposite-particle position {x, y, z} in Bohr + /// @param type "elec" or "hole" + void plot_conditional_density(const int istate, const std::array& r_fix, const std::string& type); + + /// @brief Compute conditional electron or hole density on a 2D BvK cross-section + /// Evaluates psi_cond(r) on a regular grid in the chosen plane (ab, bc, or ca), + /// spanning scale × BvK supercell extent centered on the home cell. Uses full + /// Bloch-summed orbital evaluation with home-cell caching for efficiency. + /// Writes a data file with grid + metadata readable by plot_cond_silce.py. + /// @param istate BSE state index + /// @param r_fix fixed opposite-particle position {x, y, z} in Bohr + /// @param plane cross-section plane: "ab", "bc", or "ca" + /// @param slice_pos offset along the perpendicular direction (Bohr) + /// @param npoints desired grid resolution (pts per cell ≈ npoints / total_cells) + /// @param scale view range relative to BvK supercell (default 1.3) + /// @param type "elec" or "hole" + void plot_cond_slice(const int istate, + const std::array& r_fix, + const std::string& plane, + double slice_pos, + int npoints, + double scale, + const std::string& type); + + private: + /// @brief Build one effective AO coefficient vector per k-point for a + /// conditional electron or hole wavefunction. + std::vector>> build_conditional_coefficients(int istate, + const std::array& r_fix, + bool plot_electron); + + const int nspin_x = 1; ///< 1 for singlet/triplet, 2 for up/down (open-shell) + const int naos; ///< number of atomic-orbital basis functions + const std::vector& nocc; ///< number of occupied bands per spin + const std::vector& nvirt; ///< number of virtual bands per spin + const int nk; ///< number of k-points + const int ldim; ///< local leading dimension of X (data size per state) + const double* eig; ///< excitation energies (Ry) + const T* X; ///< BSE eigenvector amplitudes + const K_Vectors& kv; ///< k-point list + std::vector> psi_ks_vec; ///< KS wavefunction per spin + const UnitCell& ucell; ///< unit cell + const std::vector& orb_cutoff_; ///< orbital cutoff radii + const Grid_Driver& gd_; ///< grid driver for neighbor search + const Parallel_Grid& Pgrid; ///< parallel real-space grid + const ModulePW::PW_Basis& rho_basis; ///< plane-wave basis for density grid + const std::string output_dir_; ///< output directory for density files + + const Parallel_2D& pc; ///< KS-orbital distribution reserved for MPI plotting + const Parallel_Orbitals& pmat; ///< parallel distribution for AO matrices + const LCAO_Orbitals* orb_ = nullptr; ///< orbital data for single-point evaluation + OrbitalEvaluator orb_eval_; ///< single-point NAO evaluator (initialized if orb_ != nullptr) +}; + +} // namespace LR_Util diff --git a/source/source_lcao/module_lr/utils/lr_io.cpp b/source/source_lcao/module_lr/utils/lr_io.cpp new file mode 100644 index 0000000000..190b84fc4f --- /dev/null +++ b/source/source_lcao/module_lr/utils/lr_io.cpp @@ -0,0 +1,948 @@ +#include "lr_io.h" +#include +#include +#include +#include +#include "source_lcao/module_lr/utils/lr_util.h" +#include +#include +#include +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif +namespace LR_IO { + const std::string FILE_COARSE = "stru_out"; + const std::string FILE_FINE_UNIFORM = "band_kpath_info"; + const std::string FILE_FINE_NONUNIFORM = "KPT_bse"; + const std::string FILE_BAND_OUT = "band_out"; + const std::string FILE_BAND_KPATH = "band_kpath_info"; + +void parse_band_out_file(int& nbands_file, int& nk_file, int& nspin_file, int& nocc_file) +{ + std::string file = PARAM.globalv.global_readin_dir + FILE_BAND_OUT; + std::ifstream ifs(file); + if (!ifs) throw std::runtime_error(file + " not found"); + std::string tmp, line; + double occ; + int nocc_count = 0; + + ifs >> nk_file >> nspin_file >> nbands_file; + std::cout << "band_out: nk = " << nk_file << std::endl; + std::cout << "band_out: nspin = " << nspin_file << std::endl; + std::cout << "band_out: nbands = " << nbands_file << std::endl; + for (int i = 0; i < 4; ++i) {std::getline(ifs, tmp); } //skip 4 lines + + while (ifs.peek() != EOF) + { + std::getline(ifs, line); + std::istringstream iss(line); + + iss >> tmp >> occ; + if (occ > 0.1) nocc_count++; + else if (occ < 0.1) break; + } + nocc_file = nocc_count; + std::cout << "nocc in band_out: " << nocc_file << std::endl; + + if (PARAM.inp.bse_use_fine_kgrid) + { + file = PARAM.globalv.global_readin_dir + FILE_BAND_KPATH; + std::ifstream ifs(file); + if (!ifs) throw std::runtime_error(file + " not found"); + std::string tmp; + ifs >> tmp >> nbands_file >> nspin_file >> nk_file; + std::cout << "band_kpath_info: nk = " << nk_file << std::endl; + std::cout << "band_kpath_info: nspin = " << nspin_file << std::endl; + std::cout << "band_kpath_info: nbands = " << nbands_file << std::endl; + std::cout << "BSE will use fine kgrid." << std::endl; + ifs.close(); + } +} + +#ifdef __EXX + +RI_kRlist::RI_kRlist(const UnitCell& ucell, + K_Vectors* const pkv) + : klist(pkv) +{ + read_kpts_coarse(PARAM.globalv.global_readin_dir + FILE_COARSE, ucell, this->klist); + this->klist_coarse = *this->klist; + this->period = RI_Util::get_Born_vonKarmen_period(*klist); + this->Rlist = RI_Util::get_Born_von_Karmen_cells(period); + // std::cout << "Rlist:" << std::endl; + // int count = 0; + // for (const auto& iR: Rlist) + // { + // count++; + // std::cout << "iR=" << count <<": "<< iR[0] << " " << iR[1] << " " << iR[2] << std::endl; + // } + if (PARAM.inp.bse_use_fine_kgrid==1) + { + read_kpts_fine(PARAM.globalv.global_readin_dir + FILE_FINE_UNIFORM, ucell, this->klist, false); + } + else if (PARAM.inp.bse_use_fine_kgrid==2) + { + read_kpts_fine(PARAM.globalv.global_readin_dir + FILE_FINE_NONUNIFORM, ucell, this->klist, true); + } + else if (PARAM.inp.bse_use_fine_kgrid!=0) + ModuleBase::WARNING_QUIT("LR_IO", "bse_use_fine_kgrid must be 0, 1 or 2"); +}; + +void RI_kRlist::read_kpts_coarse(const std::string& file, const UnitCell& ucell, K_Vectors* const klist) +{ + std::ifstream ifs; + ifs.open(file); + if (!ifs) throw std::runtime_error(file + " not found"); + std::string tmp; + for (int i = 0; i < 7; ++i) { std::getline(ifs, tmp); } // get the 7th line(number of atoms) + int nat = std::stoi(tmp); + for (int i = 0; i != nat; ++i) { std::getline(ifs, tmp); } + int nks_original = klist->get_nks(); + // std::cout << "Origianl klist (Cartesian|Direct)" << std::endl; + // for (int ik = 0;ik < nks_original;++ik) + // { + // std::cout << "ik=" << std::setw(5) << ik << std::setw(11) << klist->kvec_c[ik].x << std::setw(11) + // << klist->kvec_c[ik].y << std::setw(11) << klist->kvec_c[ik].z << " | " << std::setw(11) + // << klist->kvec_d[ik].x << std::setw(11) << klist->kvec_d[ik].y << std::setw(11) << klist->kvec_d[ik].z << std::endl; + // } + + ifs >> klist->nmp[0] >> klist->nmp[1] >> klist->nmp[2]; + int nk = klist->nmp[0] * klist->nmp[1] * klist->nmp[2]; + int nks = (PARAM.inp.nspin == 2) ? 2 * nk : nk; + assert(nks == nks_original); + + for (int ik = 0; ik < nk; ++ik) + { + ifs >> klist->kvec_c[ik].x >> klist->kvec_c[ik].y >> klist->kvec_c[ik].z; + klist->kvec_c[ik] /= ModuleBase::TWO_PI * ModuleBase::BOHR_TO_A; // in unit of 2pi/angstrom + klist->kvec_d[ik] = klist->kvec_c[ik] * ucell.latvec.Transpose(); + set_zero_if_close(klist->kvec_d[ik]); + klist->wk[ik] = 1.0 / double(nk); + } + if (PARAM.inp.nspin == 2) + { + for (int ik = 0; ik < nk; ++ik) + { + klist->kvec_c[ik + nk] = klist->kvec_c[ik]; + klist->kvec_d[ik + nk] = klist->kvec_d[ik]; + klist->wk[ik + nk] = klist->wk[ik]; + } + } + + std::ofstream ofs_kpts_coarse(PARAM.globalv.global_out_dir + "kpts_coarse.dat"); + ofs_kpts_coarse << "kpts_coarse:" << nk << std::setw(16) << "( Cartesian" << std::setw(36) + << "| Direct )" << std::setw(15) << "| wk (normalized as sum = nk)" << std::endl; + for (int ik = 0; ik < nks; ++ik) + { + ofs_kpts_coarse << std::setw(5) << ik << std::setw(12) << klist->kvec_c[ik].x << std::setw(12) + << klist->kvec_c[ik].y << std::setw(12) << klist->kvec_c[ik].z << " | " << std::setw(12) + << klist->kvec_d[ik].x << std::setw(12) << klist->kvec_d[ik].y << std::setw(12) << klist->kvec_d[ik].z + << " | " << klist->wk[ik]*nk << std::endl; + } + ofs_kpts_coarse.close(); +} + +void RI_kRlist::read_kpts_fine(const std::string& file, const UnitCell& ucell, K_Vectors* const klist, const bool is_weighted) +{ + // band_kpath_info format: first line: nband nbasis nspin nk, then kx ky kz per line (direct coords) + // KPT_bse format: first line = nk, then kx ky kz wk per line (direct coords, BSE weight sum=nk) + std::ifstream ifs; + ifs.open(file); + if (!ifs) throw std::runtime_error(file + " not found"); + int nk; + if (is_weighted) {ifs >> nk; ifs.ignore(2048, '\n');} + else {ifs >> nk >> nk >> nk >> nk;} + + int nks = (PARAM.inp.nspin == 2) ? 2 * nk : nk; + klist->set_nks(nks); + klist->set_nkstot(nks); + klist->set_nkstot_full(nk); + + auto klist_reset = [&klist](int kpoint_number){ + klist->kvec_c.resize(0); klist->kvec_c.resize(kpoint_number); + klist->kvec_d.resize(0); klist->kvec_d.resize(kpoint_number); + klist->wk.resize(0); klist->wk.resize(kpoint_number); + klist->isk.resize(0); + klist->ngk.resize(0); + }; + klist_reset(nks); + + for (int ik = 0; ik < nk; ++ik) + { + ifs >> klist->kvec_d[ik].x >> klist->kvec_d[ik].y >> klist->kvec_d[ik].z; + if (is_weighted) { + ifs >> klist->wk[ik]; + klist->wk[ik] /= double(nk); + } + else {klist->wk[ik] = 1.0 / double(nk);} + klist->kvec_c[ik] = klist->kvec_d[ik] * ucell.G; + set_zero_if_close(klist->kvec_c[ik]); + } + std::cout << "Read " << nk << " k-points and weights from " << file << std::endl; + if (PARAM.inp.nspin == 2) + { + for (int ik = 0; ik < nk; ++ik) + { + klist->kvec_c[ik + nk] = klist->kvec_c[ik]; + klist->kvec_d[ik + nk] = klist->kvec_d[ik]; + klist->wk[ik + nk] = klist->wk[ik]; + } + } + std::ofstream ofs_kpts_fine(PARAM.globalv.global_out_dir + "kpts_fine.dat"); + ofs_kpts_fine << "kpts_fine:" << nk << std::setw(18) << "( Cartesian" << std::setw(36) + << "| Direct )" << std::setw(15) << "| wk (normalized as sum = nk)" << std::endl; + for (int ik = 0; ik < nk; ++ik) + { + ofs_kpts_fine << std::setw(5) << ik << std::setw(12) << klist->kvec_c[ik].x << std::setw(12) + << klist->kvec_c[ik].y << std::setw(12) << klist->kvec_c[ik].z << " | " << std::setw(12) + << klist->kvec_d[ik].x << std::setw(12) << klist->kvec_d[ik].y << std::setw(12) << klist->kvec_d[ik].z + << " | " << klist->wk[ik]*nk << std::endl; + } + ofs_kpts_fine.close(); +} + +std::vector read_energy_qp(const int nocc, + const int nvirt, + int& ncore, + const int nk, + const int nspin_tmp, + const int nspin_file) +{ + const std::string file = PARAM.globalv.global_readin_dir + "energy_qp"; + std::cout << "in read_energy_qp, nbands(nocc+nvir): " << (nocc+nvirt) << std::endl; + std::vector eig_info( nspin_tmp * nk * (nocc + nvirt) * 3 ); // occ, eig_ks, eig_gw + std::ifstream ifs_gw (file); + if (!ifs_gw) throw std::runtime_error(file + " not found"); + std::string temp; + int read_ik; + double occ, eig_ks, eig_gw; + + for (int is =0; is < nspin_file; ++is){ + for (int ik = 0; ik < nk; ++ik){ + for (int i = 0;i < 2;++i) { std::getline(ifs_gw, temp); } // skip the first 2 lines + ifs_gw >> temp >> read_ik ; + // std::cout << "ik: " << ik <<" is:" << is << std::endl; + assert(ik == (read_ik-1)); + int ivirt = 0; + std::getline(ifs_gw, temp); // skip the interval line + std::getline(ifs_gw, temp); // skip the interval line + std::vector ks_temps; + std::vector gw_temps; + std::vector occ_temps; + while (ifs_gw.peek() != '-') + { + std::getline(ifs_gw, temp); + std::istringstream iss(temp); + iss >> temp >> occ >> eig_ks >> eig_gw; + ks_temps.push_back(eig_ks * 2); // Ha to Ry + gw_temps.push_back(eig_gw * 2); // Ha to Ry + occ_temps.push_back(occ); + if (occ < 0.1) { ivirt++;} + if (ivirt == nvirt) { break; } + } + ncore = gw_temps.size() - nocc - nvirt; + for (int ib = 0;ib < nocc + nvirt;++ib) + { + int ikstep = (ik + is * nk) * (nocc + nvirt); + eig_info[(ikstep + ib)*3] = occ_temps[ncore + ib]; + eig_info[(ikstep + ib)*3 + 1] = ks_temps[ncore + ib]; + eig_info[(ikstep + ib)*3 + 2] = gw_temps[ncore + ib]; + // std::cout <<"GW_info: ik=" << std::setw(5) << ik << " ib=" << std::setw(5) << ib + // << std::setw(9) << eig_info[(ikstep + ib)*3] << std::setw(11) + // << eig_info[(ikstep + ib)*3 + 1] << std::setw(11) + // << eig_info[(ikstep + ib)*3 + 2] << std::endl; //check + } + while (ifs_gw.peek() != '-' && ifs_gw.peek() != EOF) + { + std::getline(ifs_gw, temp); // skip the virtual bands to next k-point + } + } + } + if (nspin_file == 1 && nspin_tmp == 2) { + std::cout << "duplicate the spin channel since the gw file only has one spin channel" << std::endl; + int spin_block = nk * (nocc + nvirt) * 3; + assert(eig_info.size() == 2 * spin_block); + std::copy_n(eig_info.data(), spin_block, eig_info.data() + spin_block); + } + ifs_gw.close(); + std::cout << "Finish read gw, ncore=" << ncore << std::endl; + return eig_info; +} + +std::vector read_energy_qp_from_band_files(const K_Vectors& kv, + const int nocc, + const int nvirt, + int& ncore, + const int nk, + const int nspin_tmp, + const int nspin_file) +{ + const std::string ks_prefix = PARAM.globalv.global_readin_dir + "KS_band_spin_"; + const std::string gw_prefix = PARAM.globalv.global_readin_dir + "GW_band_spin_"; + std::cout << "in read_energy_qp_from_band_files, nbands(nocc+nvir): " << (nocc+nvirt) << std::endl; + std::vector eig_info( nspin_tmp * nk * (nocc + nvirt) * 3 ); // occ, eig_ks, eig_gw + for (int is =0; is < nspin_file; ++is) + { + std::string ks_file = ks_prefix + std::to_string(is + 1) + ".dat"; + std::string gw_file = gw_prefix + std::to_string(is + 1) + ".dat"; + std::ifstream ifs_ks (ks_file); + if (!ifs_ks) throw std::runtime_error(ks_file + " not found"); + std::ifstream ifs_gw (gw_file); + if (!ifs_gw) throw std::runtime_error(gw_file + " not found"); + std::string temp; + int read_ik; + double occ, eig_ks, eig_gw; + double kx, ky, kz; + for (int ik = 0; ik < nk; ++ik) + { + ifs_ks >> read_ik >> kx >> ky >> kz; + // std::cout << "ik: " << ik <<" is:" << is << std::endl; + double thread = 1.e-6; + assert(ik == (read_ik-1)); + assert(std::abs(kx - kv.kvec_d[ik + is * nk].x) < thread); + assert(std::abs(ky - kv.kvec_d[ik + is * nk].y) < thread); + assert(std::abs(kz - kv.kvec_d[ik + is * nk].z) < thread); + ifs_gw >> read_ik >> kx >> ky >> kz; + assert(ik == (read_ik-1)); + assert(std::abs(kx - kv.kvec_d[ik + is * nk].x) < thread); + assert(std::abs(ky - kv.kvec_d[ik + is * nk].y) < thread); + assert(std::abs(kz - kv.kvec_d[ik + is * nk].z) < thread); + int ivirt = 0; + std::vector ks_temps; + std::vector gw_temps; + std::vector occ_temps; + + std::getline(ifs_ks, temp); + std::istringstream ks_line(temp); + std::getline(ifs_gw, temp); + std::istringstream gw_line(temp); + while (true) + { + ks_line >> occ >> eig_ks; + gw_line >> occ >> eig_gw; + ks_temps.push_back(eig_ks / ModuleBase::Ry_to_eV); + gw_temps.push_back(eig_gw / ModuleBase::Ry_to_eV); + occ_temps.push_back(occ); + if (occ*nk < 0.1) { ivirt++;} + if (ivirt == nvirt) { break; } + } + + ncore = gw_temps.size() - nocc - nvirt; + for (int ib = 0;ib < nocc + nvirt;++ib) + { + int ikstep = (ik + is * nk) * (nocc + nvirt); + eig_info[(ikstep + ib)*3] = occ_temps[ncore + ib]; + eig_info[(ikstep + ib)*3 + 1] = ks_temps[ncore + ib]; + eig_info[(ikstep + ib)*3 + 2] = gw_temps[ncore + ib]; + // std::cout <<"GW_info: ib=" << std::setw(5) << ib + // << std::setw(9) << eig_info[(ikstep + ib)*3] << std::setw(11) + // << eig_info[(ikstep + ib)*3 + 1] << std::setw(11) + // << eig_info[(ikstep + ib)*3 + 2] << std::endl; //check + } + } + ifs_ks.close(); + ifs_gw.close(); + } + if (nspin_file == 1 && nspin_tmp == 2) { + std::cout << "duplicate the spin channel since the gw file only has one spin channel" << std::endl; + int spin_block = nk * (nocc + nvirt) * 3; + assert(eig_info.size() == 2 * spin_block); + std::copy_n(eig_info.data(), spin_block, eig_info.data() + spin_block); + } + std::cout << "Finish read gw, ncore=" << ncore << std::endl; + return eig_info; +} + +template +void read_librpa_eigenvectors(psi::Psi& wfc_ks, + psi::Psi& wfc_ks_global, + const std::string& path, + const int ncore, + const int nbands_file, + const int nspin_tmp, + const int nspin_file, + Parallel_Orbitals& pmat) +{ + int nbands = pmat.get_wfc_global_nbands();// nbands = nocc + nvirt + int nbasis = pmat.get_wfc_global_nbasis(); + assert(nbands == wfc_ks_global.get_nbands()); + assert(nbasis == wfc_ks_global.get_nbasis()); + assert((ncore + nbands) <= nbands_file); + const size_t nk = PARAM.inp.nspin == 2 ? wfc_ks.get_nk() / 2 : wfc_ks.get_nk(); + + if (GlobalV::MY_RANK == 0) + { + struct dirent *ptr; + DIR *dir; + dir = opendir(path.c_str()); + std::vector readen_k(nk, false); + + while ((ptr = readdir(dir)) != NULL){// read all the files in the directory + std::string fm(ptr->d_name); + if (fm.find("KS_eigenvector") == 0)// find file KS_eigenvectorXXX + { + //std::cout << "found librpa_eigenvector file:" << fm << std::endl; + std::ifstream file_librpa_ks(path + fm); + std::string tmp; + while (file_librpa_ks.peek() != EOF) + { + int ik; + file_librpa_ks >> ik; + ik = ik - 1; + assert(readen_k[ik] == false); + file_librpa_ks >> std::ws; // skip the blank and '\n' to get the next content + for (int iw = 0; iw < nbasis; ++iw) { + for (int ib = 0; ib < nbands_file; ++ib) { + for (int is = 0; is < nspin_file; ++is) { + if (ib >= ncore && ib< (ncore+nbands)) { + LR_IO::read_one_data(file_librpa_ks, wfc_ks_global(ik+is*nk, ib-ncore, iw)); + file_librpa_ks >> std::ws; + } + else { + std::getline(file_librpa_ks, tmp); //skip the useless bands + } + } + } + } + if (nspin_tmp == 2 && nspin_file == 1) { + std::copy_n(&wfc_ks_global(ik,0,0), nbands*nbasis, &wfc_ks_global(ik + nk,0,0)); + } + readen_k[ik] = true; + } + file_librpa_ks.close(); + } + } + closedir(dir); + for(int ik = 0; ik < nk; ++ik) { + if (!readen_k[ik]) + throw std::runtime_error("librpa_eigenvector file not found for k-point " + std::to_string(ik+1)); + } + + }// end of if (GlobalV::MY_RANK == 0); next MPI_comm to other ranks + + // change wfc_ks_global phase to make arg() = 0 + std::cout << "Do phase correction" << std::endl; + for (int iks = 0; iks < wfc_ks.get_nk(); ++iks){ + if (GlobalV::MY_RANK == 0) { + // test: output wfc + // std::cout << "wfc_gs_read_from_librpa for iks:" << iks << std::endl; + // for (int ib = 0; ib < nbands; ++ib) + // { + // std::cout << "band " << ib << ": "; + // for (int iw = 0; iw < nbasis; ++iw) + // { + // std::cout << wfc_ks_global(iks, ib, iw) << " "; + // } + // std::cout << std::endl; + // } + if (iks != 0) + { + for(int ib = 0; ib < nbands; ++ib) + { + TK phase = LR_Util::inner_product(&wfc_ks_global(iks,ib,0), &wfc_ks_global(0,ib,0), nbasis); + phase = phase / std::abs(phase); + for (int iw = 0; iw < nbasis; ++iw) + { + wfc_ks_global(iks, ib, iw) *= phase; + } + // TK test_phase = LR_Util::inner_product(&wfc_ks_global(iks,ib,0), &wfc_ks_global(0,ib,0), nbasis); + // std::cout << "After phase correction, iks, ib, phase: " << iks << " " << ib << " " << test_phase << std::endl; + } + } + } + + wfc_ks_global.fix_k(iks); + wfc_ks.fix_k(iks); +#ifdef __MPI + MPI_Bcast(wfc_ks_global.get_pointer(), nbands * nbasis, LR_Util::MPIType::value(), 0, MPI_COMM_WORLD); + Parallel_2D pv_glb; + pv_glb.set(nbasis, nbands, std::max(nbasis, nbands), pmat.blacs_ctxt); + Cpxgemr2d(nbasis, nbands, wfc_ks_global.get_pointer(), 1, 1, pv_glb.desc, + wfc_ks.get_pointer(), 1, 1, const_cast(pmat.desc_wfc)/*nbasis×nbands*/, + pv_glb.blacs_ctxt); +#else + BlasConnector::copy(nbands*nlocal, wfc_ks_global.get_pointer(), 1, wfc_ks.get_pointer(), 1); +#endif + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read librpa eigenvectors."); +} + +template +void read_librpa_eigenvectors_from_band_files(psi::Psi& wfc_ks, + psi::Psi& wfc_ks_global, + const std::string& path, + const int ncore, + const int nbands_file, + const int nspin_tmp, + const int nspin_file, + Parallel_Orbitals& pmat) +{ + int nbands = pmat.get_wfc_global_nbands();// nbands = nocc + nvirt + int nbasis = pmat.get_wfc_global_nbasis(); + assert(nbands == wfc_ks_global.get_nbands()); + assert(nbasis == wfc_ks_global.get_nbasis()); + assert((ncore + nbands) <= nbands_file); + const size_t nk = PARAM.inp.nspin == 2 ? wfc_ks.get_nk() / 2 : wfc_ks.get_nk(); + + if (GlobalV::MY_RANK == 0) + { + for (int ik = 0; ik < nk; ++ik) + { + std::stringstream ss; + ss << path << "band_KS_eigenvector_k_" << std::setfill('0') << std::setw(5) << ik + 1 << ".txt"; + std::ifstream infile(ss.str(), std::ios::binary); + if (!infile) + { + throw std::runtime_error("Error: Cannot open file " + ss.str()); + } + size_t total_complex = static_cast(nspin_file * nbands_file * nbasis); + size_t total_doubles = total_complex * 2; + + std::vector double_buffer(total_doubles); + infile.read(reinterpret_cast(double_buffer.data()), total_doubles * sizeof(double)); + if (infile.gcount() != static_cast(total_doubles * sizeof(double))) + { + throw std::runtime_error("Error: failed to read " + ss.str()); + } + for (int is = 0; is < nspin_file; ++is) { + for (int ib = ncore; ib < (ncore+nbands); ++ib) { + for (int iw = 0; iw < nbasis; ++iw) { + int index = 2 * (is * nbands_file * nbasis + ib * nbasis + iw); + LR_IO::read_one_data(double_buffer[index], double_buffer[index+1], wfc_ks_global(ik+is*nk, ib-ncore, iw)); + } + } + } + if (nspin_tmp == 2 && nspin_file == 1) { + std::copy_n(&wfc_ks_global(ik,0,0), nbands*nbasis, &wfc_ks_global(ik + nk,0,0)); + } + infile.close(); + } + }// end of if (GlobalV::MY_RANK == 0); next MPI_comm to other ranks + + // change wfc_ks_global phase to make arg() = 0 + std::cout << "Do phase correction" << std::endl; + for (int iks = 0; iks < wfc_ks.get_nk(); ++iks){ + if (GlobalV::MY_RANK == 0) { + // test: output wfc + // std::cout << "wfc_gs_read_from_librpa for iks:" << iks << std::endl; + // for (int ib = 0; ib < nbands; ++ib) + // { + // std::cout << "band " << ib << ": "; + // for (int iw = 0; iw < nbasis; ++iw) + // { + // std::cout << wfc_ks_global(iks, ib, iw) << " "; + // } + // std::cout << std::endl; + // } + if (iks != 0) + { + for(int ib = 0; ib < nbands; ++ib) + { + TK phase = LR_Util::inner_product(&wfc_ks_global(iks,ib,0), &wfc_ks_global(0,ib,0), nbasis); + phase = phase / std::abs(phase); + for (int iw = 0; iw < nbasis; ++iw) + { + wfc_ks_global(iks, ib, iw) *= phase; + } + // TK test_phase = LR_Util::inner_product(&wfc_ks_global(iks,ib,0), &wfc_ks_global(0,ib,0), nbasis); + // std::cout << "After phase correction, iks, ib, phase: " << iks << " " << ib << " " << test_phase << std::endl; + } + } + } + + wfc_ks_global.fix_k(iks); + wfc_ks.fix_k(iks); +#ifdef __MPI + MPI_Bcast(wfc_ks_global.get_pointer(), nbands * nbasis, LR_Util::MPIType::value(), 0, MPI_COMM_WORLD); + Parallel_2D pv_glb; + pv_glb.set(nbasis, nbands, std::max(nbasis, nbands), pmat.blacs_ctxt); + Cpxgemr2d(nbasis, nbands, wfc_ks_global.get_pointer(), 1, 1, pv_glb.desc, + wfc_ks.get_pointer(), 1, 1, const_cast(pmat.desc_wfc)/*nbasis×nbands*/, + pv_glb.blacs_ctxt); +#else + BlasConnector::copy(nbands*nlocal, wfc_ks_global.get_pointer(), 1, wfc_ks.get_pointer(), 1); +#endif + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read librpa eigenvectors."); +} + +template // only for blocking by atom pairs +TLRI read_coulomb_mat_k(const std::string& path, const TLRI& Cs, LR_IO::RI_kRlist& kRlist) +{ + struct dirent *ptr; + DIR *dir; + bool has_unshrinked = false; + dir = opendir(path.c_str()); + if (!dir) + { + throw std::runtime_error("Cannot open directory: " + path); + } + while ((ptr = readdir(dir)) != nullptr) + { + std::string fm(ptr->d_name); + if (fm.find("coulomb_unshrinked_cut_") == 0) + { + has_unshrinked = true; + break; + } + } + rewinddir(dir); + const std::string prefix = has_unshrinked ? "coulomb_unshrinked_cut_" : "coulomb_cut_"; + std::cout << "read_coulomb_mat_k: using prefix \"" << prefix << "\" in directory \"" << path << "\"" << std::endl; + + size_t nk = 0, nabf = 0, istart = 0, jstart = 0, iend = 0, jend = 0; + + K_Vectors* const klist = &(kRlist.klist_coarse); + int klist_nk = klist->nmp[0] * klist->nmp[1] * klist->nmp[2]; + int ik_readin = -1; + TLRI Vs; + std::map, RI::Tensor>>> Vq; // , T>> + const int nat = Cs.size(); + std::vector abf_start_index(nat+1, 1); + for (int iat = 0; iat < nat; ++iat) + { + abf_start_index[iat+1] = abf_start_index[iat] + Cs.at(iat).at({ 0, {0,0,0} }).shape[0]; + } + + auto to_atom = [&](const int start, const int end) -> int + { + for (int iat = 0;iat < nat;++iat) + { + size_t abf_start = abf_start_index[iat]; + size_t abf_end = abf_start_index[iat+1] - 1; + if (start == abf_start && end == abf_end) + { + return iat; + } + } + throw std::runtime_error("Error in read_coulomb_mat_k: cannot find the atom for given auxiliary basis set range"); + }; + + while ((ptr = readdir(dir)) != NULL){// read all the files in the directory + std::string fm(ptr->d_name); + if (fm.find(prefix) == 0)// find file coulomb_cut_xxx + { + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "found Coulomb file: " + fm + ", start reading..."); + std::ifstream ifs(path + fm); + ifs >> nk;// actual nk + assert(nk == klist_nk); + + while (ifs.peek() != EOF) + { + ifs >> nabf >> istart >> iend >> jstart >> jend >> ik_readin >> klist->wk[ik_readin-1]; + if (ifs.peek() == EOF) { break; } + int ik = ik_readin - 1; + int iat1 = to_atom(istart, iend); + int iat2 = to_atom(jstart, jend); + const size_t nabf1 = iend - istart + 1; + const size_t nabf2 = jend - jstart + 1; + + RI::Tensor> t({ nabf1, nabf2 }); + for (int i = 0;i < nabf1;++i) + { + for (int j = 0;j < nabf2;++j) + { + LR_IO::read_one_data(ifs, t(i, j)); + } + } + Vq[iat1][{iat2, ik}] = t; + + if (iat1 != iat2) + { // coulomb_mat has only the upper triangle part + Vq[iat2][{iat1, ik}] = t.dagger(); + continue; + } + } + ifs.close(); + } + } + closedir(dir); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read Vq files. Now convert Vq to VR."); + for (const TC& R : kRlist.Rlist ) + { + for (int iat1 = 0;iat1 < nat;++iat1) + { + for (int iat2 = 0;iat2 < nat;++iat2) + { + Vs[iat1][{iat2, R}] = RI::Tensor(Vq[iat1][{iat2, 0}].shape); + } + } + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "VR keys has been prepared."); +#ifdef _OPENMP +#pragma omp parallel for schedule(dynamic) collapse(3) +#endif + for (const TC& R : kRlist.Rlist ) + { +// #ifdef _OPENMP +// #pragma omp critical +// std::cout<<"thread"<< omp_get_thread_num() << " convert V: R="<(Vq[iat1][{iat2, 0}].shape); + for (int ik = 0;ik < nk;++ik) + { + const ModuleBase::Vector3& kvec = klist->kvec_d.at(ik); + const double arg = -1.0 * ModuleBase::TWO_PI * (kvec.x * R[0] + kvec.y * R[1] + kvec.z * R[2]); + const std::complex kphase (cos(arg), sin(arg)); + Vs[iat1][{iat2, R}] += RI::Global_Func::convert (Vq[iat1][{iat2, ik}] * kphase) * RI::Global_Func::convert(klist->wk[ik]); + } + } + } + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "convert Vq to VR."); + ModuleBase::TITLE("LR_IO", "read_Vs done."); + return Vs; +} + +template // any blocking +TLRI read_coulomb_mat_general_k(const std::string& path, const TLRI& Cs, LR_IO::RI_kRlist& kRlist) +{ + struct dirent *ptr; + DIR *dir; + bool has_unshrinked = false; + dir = opendir(path.c_str()); + if (!dir) + { + throw std::runtime_error("Cannot open directory: " + path); + } + while ((ptr = readdir(dir)) != nullptr) + { + std::string fm(ptr->d_name); + if (fm.find("coulomb_unshrinked_cut_") == 0) + { + has_unshrinked = true; + break; + } + } + rewinddir(dir); + const std::string prefix = has_unshrinked ? "coulomb_unshrinked_cut_" : "coulomb_cut_"; + std::cout << "read_coulomb_mat_k: using prefix \"" << prefix << "\" in directory \"" << path << "\"" << std::endl; + + TLRI Vs; + std::map, RI::Tensor>>> Vq; // , T>> + std::map>> Vq_tmp; // + + size_t nabf = 0, istart = 0, jstart = 0, iend = 0, jend = 0; + int nk = 0; + K_Vectors* const klist = &(kRlist.klist_coarse); + int klist_nk = klist->nmp[0] * klist->nmp[1] * klist->nmp[2]; + int ik_readin = -1; + + while ((ptr = readdir(dir)) != NULL){// read all the files in the directory + std::string fm(ptr->d_name); + if (fm.find(prefix) == 0)// find file coulomb_cut_xxx + { + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "found Coulomb file: " + fm + ", start reading..."); + std::ifstream ifs(path + fm); + ifs >> nk; // actual nk + assert(nk == klist_nk); + + while (ifs.peek() != EOF) + { + ifs >> nabf >> istart >> iend >> jstart >> jend >> ik_readin >> klist->wk[ik_readin-1];// wk is not used + if (ifs.peek() == EOF) { break; } + int ik = ik_readin - 1; + if (Vq_tmp[ik].empty()) { Vq_tmp[ik].resize(nabf * nabf, 0.0); } + auto& Vq_tmp_k = Vq_tmp.at(ik); + for (int i = istart - 1;i < iend;++i) + { + for (int j = jstart - 1;j < jend;++j) + { + LR_IO::read_one_data(ifs, Vq_tmp_k[i * nabf + j]); + } + } + } + ifs.close(); + } + } + closedir(dir); + + const int nat = Cs.size(); + istart = 0; + for (int iat1 = 0;iat1 < nat;++iat1) + { + const size_t nabf1 = Cs.at(iat1).at({ 0, {0,0,0} }).shape[0]; + jstart = 0; + for (int iat2 = 0;iat2 < nat;++iat2) + { + const size_t nabf2 = Cs.at(iat2).at({ 0, {0,0,0} }).shape[0]; + for (int ik = 0; ik < nk; ++ik){ + if (iat1 > iat2) + { // coulomb_mat has only the upper triangle part + Vq[iat1][{iat2, ik}] = Vq[iat2][{iat1, ik}].dagger(); + } + else + { + RI::Tensor> t({ nabf1, nabf2 }); + for (int i = 0;i < nabf1;++i) + { + for (int j = 0;j < nabf2;++j) + { + t(i, j) = Vq_tmp[ik][(istart + i) * nabf + jstart + j]; + } + } + Vq[iat1][{iat2, ik}] = t; + } + } + jstart += nabf2; + } + assert(jstart == nabf); + istart += nabf1; + } + assert(istart == nabf); + + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read Vq files. Now convert Vq to VR."); + for (const TC& R : kRlist.Rlist ) + { + for (int iat1 = 0;iat1 < nat;++iat1) + { + for (int iat2 = 0;iat2 < nat;++iat2) + { + Vs[iat1][{iat2, R}] = RI::Tensor({ Vq[iat1][{iat2, 0}].shape[0], Vq[iat1][{iat2, 0}].shape[1] }); + } + } + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "VR keys has been prepared."); + + double reciprocal_nk = 1.0 / double(nk); +#ifdef _OPENMP +#pragma omp parallel for schedule(dynamic) collapse(3) +#endif + for (const TC& R : kRlist.Rlist ) + { +// #ifdef _OPENMP +// #pragma omp critical +// std::cout<<"thread"<< omp_get_thread_num() << "convert V: R="<& kvec = klist->kvec_d.at(ik); + const double arg = -1.0 * ModuleBase::TWO_PI * (kvec.x * R[0] + kvec.y * R[1] + kvec.z * R[2]); + const std::complex kphase (cos(arg), sin(arg)); + Vs[iat1][{iat2, R}] += RI::Global_Func::convert (Vq[iat1][{iat2, ik}] * kphase) * RI::Global_Func::convert(reciprocal_nk); + } + } + } + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "convert Vq to VR."); + ModuleBase::TITLE("LR_IO", "read_Vs done."); + return Vs; +} + +template +TLRI read_Ws(const TLRI& Vs, const std::vector& Rlist) +{ + ModuleBase::TITLE("LR_IO", "read_Ws"); + std::map>> Ws; + + const int nat = Vs.size(); + std::string temp; + int nk, istart, iend, jstart, jend, ik; + size_t nabfmu, nabfnu, non_zero, mu, nu; //I.nab, J.nab + size_t nR = Rlist.size(); + for(int iat = 0; iat != nat; ++iat)//loop atom I + { + for(int jat = 0; jat != nat; ++jat)//loop atom J + { + for(int iR = 0; iR < nR; ++iR) + { + std::ifstream infileW; + std::string filename = "librpa.d/Wc_Mu_"+std::to_string(iat)+"_Nu_"+std::to_string(jat)+"_iR_"+std::to_string(iR)+"_ifreq_0.mtx"; + infileW.open(filename); + if(!infileW) throw std::runtime_error( filename + " not found!"); + // else std::cout << "reading Wc file: " << filename ; + int nabf1 = Vs.at(iat).at({jat,{0,0,0}}).shape[0]; + int nabf2 = Vs.at(iat).at({jat,{0,0,0}}).shape[1]; + + TC R; // iR of Wc file is not equal to iR in Rlist !!! + infileW.ignore(2048, '\n'); // skip line 1: %%MatrixMarket... + std::getline(infileW, temp); // read line 2: "%" + std::getline(infileW, temp); // read line 3: "% Wc at iR N ( Rx Ry Rz ) ..." + size_t lparen = temp.find('('); + size_t rparen = temp.find(')', lparen); + if (lparen == std::string::npos || rparen == std::string::npos) + throw std::runtime_error("Failed to parse R coordinates in " + filename); + std::istringstream riss(temp.substr(lparen + 1, rparen - lparen - 1)); + riss >> R[0] >> R[1] >> R[2]; + while(infileW.peek() == '%') infileW.ignore(2048, '\n'); //skip comments + + infileW >> nabfmu >> nabfnu >> non_zero; + assert(nabfmu == nabf1); + assert(nabfnu == nabf2); + RI::Tensor tensor_W({ nabfmu, nabfnu }); + for (int index = 0; index < non_zero; ++index) + { + infileW >> mu >> nu ; + LR_IO::read_one_data(infileW, tensor_W(mu-1, nu-1)); + } + infileW.close(); + tensor_W += Vs.at(iat).at({jat, R}); + // for(int i = 0; i != nabf1; ++i) + // for(int j = 0; j != nabf2; ++j) + // { + // tensor_W(i, j) += Vs.at(iat).at({jat, R})(i,j); + // std::cout << "Wxc: " << i << " " << j << " " << tensor_W(i,j) << std::endl; //check + // } + Ws[iat][{jat, R}] = std::move(tensor_W); + // std::cout << " Finished. R: " << "( " << R[0] << " " << R[1] << " " << R[2] << " )" << std::endl; + } + } + } + ModuleBase::TITLE("LR_IO", "read_Ws done."); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "read WR files."); + return Ws; +} + +template void read_librpa_eigenvectors( + psi::Psi& wfc_ks, psi::Psi& wfc_ks_global, + const std::string& path, const int ncore, const int nbands_file, + const int nspin_tmp, const int nspin_file, Parallel_Orbitals& pmat); + +template void read_librpa_eigenvectors>( + psi::Psi>& wfc_ks, psi::Psi>& wfc_ks_global, + const std::string& path, const int ncore, const int nbands_file, + const int nspin_tmp, const int nspin_file, Parallel_Orbitals& pmat); + +template void read_librpa_eigenvectors_from_band_files( + psi::Psi& wfc_ks, psi::Psi& wfc_ks_global, + const std::string& path, const int ncore, const int nbands_file, + const int nspin_tmp, const int nspin_file, Parallel_Orbitals& pmat); + +template void read_librpa_eigenvectors_from_band_files>( + psi::Psi>& wfc_ks, psi::Psi>& wfc_ks_global, + const std::string& path, const int ncore, const int nbands_file, + const int nspin_tmp, const int nspin_file, Parallel_Orbitals& pmat); + +template TLRI read_coulomb_mat_k +(const std::string& path, const TLRI& Cs, LR_IO::RI_kRlist& kRlist); + +template TLRI> read_coulomb_mat_k, std::complex> +(const std::string& path, const TLRI>& Cs, LR_IO::RI_kRlist& kRlist); + +template TLRI read_coulomb_mat_general_k +(const std::string& path, const TLRI& Cs, LR_IO::RI_kRlist& kRlist); + +template TLRI> read_coulomb_mat_general_k, std::complex> +(const std::string& path, const TLRI>& Cs, LR_IO::RI_kRlist& kRlist); + +template TLRI read_Ws +(const TLRI& Vs, const std::vector& Rlist); + +template TLRI> read_Ws, std::complex> +(const TLRI>& Vs, const std::vector& Rlist); + +#endif // __EXX + +} // namespace LR_IO diff --git a/source/source_lcao/module_lr/utils/lr_io.h b/source/source_lcao/module_lr/utils/lr_io.h new file mode 100644 index 0000000000..a8435ff0c3 --- /dev/null +++ b/source/source_lcao/module_lr/utils/lr_io.h @@ -0,0 +1,122 @@ +#pragma once +#include "source_base/tool_title.h" +#include "source_basis/module_ao/parallel_orbitals.h" +#include "source_cell/klist.h" +#include "source_cell/unitcell.h" +#include "source_io/module_parameter/parameter.h" +#include "source_psi/psi.h" +#ifdef __EXX +#include "source_lcao/module_ri/RI_Util.h" // for get_Born_von_Karmen_cells +#include +#endif +#include +#include +#include +#include +#include +#include + +namespace LR_IO +{ + +inline void read_one_data(std::ifstream& ifs, double& data) +{ + std::string temp; + ifs >> data >> temp; +} +inline void read_one_data(std::ifstream& ifs, std::complex& data) +{ + double real, imag; + ifs >> real >> imag; + data = std::complex(real, imag); +} +inline void read_one_data(double& re, double& im, std::complex& out) +{ + out = std::complex(re, im); +} +inline void read_one_data(double& re, double& im, double& out) +{ + out = re; +} +inline void set_zero_if_close(ModuleBase::Vector3& vec, const double tol=1e-8) +{ + if (std::abs(vec.x) < tol) vec.x = 0.0; + if (std::abs(vec.y) < tol) vec.y = 0.0; + if (std::abs(vec.z) < tol) vec.z = 0.0; +} + +void parse_band_out_file(int& nbands_file, int& nk_file, int& nspin_file, int& nocc_file); + +#ifdef __EXX +using TA = int; +using TC = std::array; +using TAC = std::pair; +template +using TLRI = std::map>>; + +class RI_kRlist +{ + public: + K_Vectors* klist = nullptr; // store fine kgrid if bse_use_fine_kgrid + K_Vectors klist_coarse; + TC period; + std::vector Rlist; + RI_kRlist() = default; + RI_kRlist(const UnitCell& ucell, K_Vectors* pkv); + ~RI_kRlist() = default; + void read_kpts_coarse(const std::string& file, const UnitCell& ucell, K_Vectors* const klist); + void read_kpts_fine(const std::string& file, const UnitCell& ucell, K_Vectors* const klist, const bool is_weighted); +}; + +/// @brief vector as {ik, iband, } +/// @param ncore: as output, number of core orbitals parsed from file +std::vector read_energy_qp(const int nocc, + const int nvirt, + int& ncore, + const int nk, + const int nspin_tmp, + const int nspin_file); + +/// @brief same as read_energy_qp, but read from KS_band and GW_band files +std::vector read_energy_qp_from_band_files(const K_Vectors& kv, + const int nocc, + const int nvirt, + int& ncore, + const int nk, + const int nspin_tmp, + const int nspin_file); + +template +void read_librpa_eigenvectors(psi::Psi& wfc_ks, + psi::Psi& wfc_ks_global, + const std::string& path, + const int ncore, + const int nbands_file, + const int nspin_tmp, + const int nspin_file, + Parallel_Orbitals& pmat); + +template +void read_librpa_eigenvectors_from_band_files(psi::Psi& wfc_ks, + psi::Psi& wfc_ks_global, + const std::string& path, + const int ncore, + const int nbands_file, + const int nspin_tmp, + const int nspin_file, + Parallel_Orbitals& pmat); + +/// only for blocking by atom pairs (abacus type) +template +TLRI read_coulomb_mat_k(const std::string& path, const TLRI& Cs, LR_IO::RI_kRlist& kRlist); + +/// for any way of blocking (aims type) +template +TLRI read_coulomb_mat_general_k(const std::string& path, const TLRI& Cs, LR_IO::RI_kRlist& kRlist); + +/// @brief read Wxc(R) = Wc(R) + Vx(R) from file +template +std::map>> read_Ws(const TLRI& Vs, const std::vector& Rlist); +#endif + +} // namespace LR_IO \ No newline at end of file diff --git a/source/source_lcao/module_lr/utils/lr_util.h b/source/source_lcao/module_lr/utils/lr_util.h index f2764fe179..a1b2e60c06 100644 --- a/source/source_lcao/module_lr/utils/lr_util.h +++ b/source/source_lcao/module_lr/utils/lr_util.h @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include #include "source_base/matrix.h" @@ -47,6 +48,28 @@ namespace LR_Util // Operators to calculate xc kernel have been moved into lr_util_xc.hpp. /// =================ALGORITHM==================== + inline double inner_product(const double* vec1, const double* vec2, const int& size) + { + double result = 0.0; + for (int i = 0; i < size; ++i) + { + result += vec1[i] * vec2[i]; + } + return result; + } + + inline std::complex inner_product(const std::complex* vec1, + const std::complex* vec2, + const int& size) + { + std::complex result = 0.0; + for (int i = 0; i < size; ++i) + { + result += std::conj(vec1[i]) * vec2[i]; + } + return result; + } + //====== newers and deleters======== /// @brief delete 2d pointer template @@ -73,6 +96,11 @@ namespace LR_Util template void matsym(T* inout, const int n, const Parallel_2D& pmat); #endif + template + bool is_hermitian(const T* mat, const Parallel_2D& pmat, const double threshold); + + template + bool is_symmetric(const T* mat, const Parallel_2D& pmat, const double threshold); ///===================Psi wrapper================= /// psi(nk=1, nbands=nb, nk * nbasis) -> psi(nb, nk, nbasis) without memory copy @@ -85,14 +113,46 @@ namespace LR_Util ///=================2D-block Parallel=============== // pack the process to setup 2d divion reusing blacs_ctxt of a new 2d-matrix void setup_2d_division(Parallel_2D& pv, int nb, int gr, int gc); - + + /// @brief assign global X to 2d-matrix, its col is band(excition state), and row is { spin, k-point, occ, virt } + /// @attention pX is 2d-blocked as {occ, virt}, this assignment is used to calculate transition density matrix c_b X_{bj} c_j + template + void global2local_X(T* local_X, const T* global_X, const int nband, const int nk, + const std::vector& nocc, const std::vector& nvirt, const std::vector& pX, + const bool openshell); #ifdef __MPI // pack the process to setup 2d divion reusing blacs_ctxt of an existing 2d-matrix void setup_2d_division(Parallel_2D& pv, int nb, int gr, int gc, const int& blacs_ctxt_in); + + /// @brief Struct to get MPI_traits for different data types + template + struct MPIType { + static MPI_Datatype value() { return MPI_DATATYPE_NULL; } + }; + template <> + struct MPIType { + static MPI_Datatype value() { return MPI_INT; } + }; + template <> + struct MPIType { + static MPI_Datatype value() { return MPI_DOUBLE; } + }; + template <> + struct MPIType> { + static MPI_Datatype value() { return MPI_DOUBLE_COMPLEX; } + }; + /// @brief assign X in pA to X in pX + template + void pA2pX(T* X_pX, const T* X_pA, const int nband, const int nk, + const std::vector& nocc, const std::vector& nvirt, + const std::vector& pX, const Parallel_2D& pA, + const int row_offset, const int col_offset, const bool openshell); + /// @brief gather 2d matrix to full matrix /// the defination of row and col is consistent with setup_2d_division template - void gather_2d_to_full(const Parallel_2D& pv, const T* submat, T* fullmat, bool col_first, int global_nrow, int global_ncol); + void gather_2d_to_full(const Parallel_2D& pv, const T* submat, T* fullmat, + const bool row_major, const std::size_t global_nrow, const std::size_t global_ncol); #endif ///=================diago-lapack==================== diff --git a/source/source_lcao/module_lr/utils/lr_util.hpp b/source/source_lcao/module_lr/utils/lr_util.hpp index 0ef1280c99..aac2979eeb 100644 --- a/source/source_lcao/module_lr/utils/lr_util.hpp +++ b/source/source_lcao/module_lr/utils/lr_util.hpp @@ -1,10 +1,13 @@ #pragma once -#include #include "lr_util.h" -#include #include "source_cell/unitcell.h" #include "source_base/constants.h" #include "source_hamilt/module_xc/xc_functional.h" +#include "source_base/module_external/lapack_connector.h" +#include "source_base/module_external/scalapack_connector.h" +#include +#include +#include namespace LR_Util { /// =================PHYSICS==================== @@ -14,7 +17,7 @@ namespace LR_Util int nelec = 0; for (int it = 0; it < ucell.ntype; ++it) { nelec += ucell.atoms[it].ncpp.zv * ucell.atoms[it].na; -} + } return nelec; } @@ -71,28 +74,154 @@ namespace LR_Util { for (int i = 0; i < n; ++i) { out[i * n + i] = 0.5 * in[i * n + i] + 0.5 * get_conj(in[i * n + i]); -} + } for (int i = 0;i < n;++i) { for (int j = i + 1;j < n;++j) { out[i * n + j] = 0.5 * (in[i * n + j] + get_conj(in[j * n + i])); out[j * n + i] = get_conj(out[i * n + j]); } -} + } } template void matsym(T* inout, const int n) { for (int i = 0; i < n; ++i) { inout[i * n + i] = 0.5 * (inout[i * n + i] + get_conj(inout[i * n + i])); -} + } for (int i = 0;i < n;++i) { for (int j = i + 1;j < n;++j) { inout[i * n + j] = 0.5 * (inout[i * n + j] + get_conj(inout[j * n + i])); inout[j * n + i] = get_conj(inout[i * n + j]); } -} + } + } + template + bool is_hermitian(const T* mat, const Parallel_2D& pmat, const double threshold){ + bool is_herm = true; + int n = pmat.get_global_row_size(); + assert(n == pmat.get_global_col_size()); + double work_dummy = 0.0; + const char norm_type = 'F'; + double norm1(0.0), norm2(0.0); +#ifdef __MPI + const int one = 1; + T alpha = 1.0; + T beta = 0.0; + if (pmat.get_local_size() > std::numeric_limits::max()) + { + std::cerr<< "Warning: pmat in RANK " + std::to_string(GlobalV::MY_RANK) + " has " + + std::to_string(pmat.get_local_size()) + " elements, which may overflow during `tranc`, please use more mpi." << std::endl; + } + std::vector herm_mat(pmat.get_local_size()); + ScalapackConnector::tranc(n, n, + alpha, const_cast(mat), one, one, pmat.desc, + beta, herm_mat.data(), one, one, pmat.desc); + ModuleBase::TITLE("LR_Util", "is_hermitian: tranc"); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "is_hermitian: tranc"); + T diff; + for (std::size_t i = 0;i < pmat.get_local_size();++i) { + diff = mat[i] - herm_mat[i]; + if (std::abs(diff) > threshold) { is_herm = false; } + norm1 += static_cast(std::norm(diff)); + norm2 += static_cast(std::norm(mat[i] + herm_mat[i])); + } + MPI_Allreduce(MPI_IN_PLACE, &norm1, 1, MPI_DOUBLE, MPI_SUM, pmat.comm()); + MPI_Allreduce(MPI_IN_PLACE, &norm2, 1, MPI_DOUBLE, MPI_SUM, pmat.comm()); + norm1 = std::sqrt(norm1); + norm2 = std::sqrt(norm2); + // We don't use `lange` here because minus_mat and sum_mat are memory-consuming. + // double norm1 = ScalapackConnector::lange(norm_type, n, n, minus_mat.data(), one, one, pmat.desc, &work_dummy); + // double norm2 = ScalapackConnector::lange(norm_type, n, n, sum_mat.data(), one, one, pmat.desc, &work_dummy); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "Frobenius norm"); + int local_flag = is_herm ? 1 : 0; + int global_flag = 0; + MPI_Allreduce(&local_flag, &global_flag, 1, MPI_INT, MPI_LAND, pmat.comm()); + is_herm = (global_flag != 0); +#else + assert(pmat.is_serial()); + std::vector minus_mat(pmat.get_local_size()); + std::vector sum_mat(pmat.get_local_size()); + for (std::size_t i = 0;i < n;++i) { + for (std::size_t j = i;j < n;++j) { + minus_mat[i * n + j] = mat[i * n + j] - get_conj(mat[j * n + i]); + minus_mat[j * n + i] = -get_conj(minus_mat[i * n + j]); + if (std::abs(minus_mat[i * n + j]) > threshold) { is_herm = false; } + sum_mat[i * n + j] = mat[i * n + j] + get_conj(mat[j * n + i]); + sum_mat[j * n + i] = get_conj(sum_mat[i * n + j]); + } + } + norm1 = LapackConnector::lange(norm_type, n, n, minus_mat.data(), n, &work_dummy); + norm2 = LapackConnector::lange(norm_type, n, n, sum_mat.data(), n, &work_dummy); +#endif + if (GlobalV::MY_RANK == 0) { + std::cout << "| Hermitian check: ||A - A^H||_F = " << norm1 << ", ||A + A^H||_F = " << norm2 << std::endl; + std::cout << "| ||A - A^H||_F / ||A + A^H||_F = " << norm1 / norm2 << std::endl; + } + return is_herm; + } + + template + bool is_symmetric(const T* mat, const Parallel_2D& pmat, const double threshold){ + bool is_sym = true; + int n = pmat.get_global_row_size(); + assert(n == pmat.get_global_col_size()); + double work_dummy = 0.0; + const char norm_type = 'F'; + double norm1(0.0), norm2(0.0); +#ifdef __MPI + const int one = 1; + T alpha = 1.0; + T beta = 0.0; + if (pmat.get_local_size() > std::numeric_limits::max()) + { + std::cerr<< "Warning: pmat in RANK " + std::to_string(GlobalV::MY_RANK) + " has " + + std::to_string(pmat.get_local_size()) + " elements, which may overflow during `tranu`, please use more mpi." << std::endl; + } + std::vector trans_mat(pmat.get_local_size()); + ScalapackConnector::tranu(n, n, + alpha, const_cast(mat), one, one, pmat.desc, + beta, trans_mat.data(), one, one, pmat.desc); + ModuleBase::TITLE("LR_Util", "is_symmetric: tranu"); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "is_hermitian: tranu"); + T diff; + for (std::size_t i = 0;i < pmat.get_local_size();++i) { + diff = mat[i] - trans_mat[i]; + if (std::abs(diff) > threshold) { is_sym = false; } + norm1 += static_cast(std::norm(diff)); + norm2 += static_cast(std::norm(mat[i] + trans_mat[i])); + } + MPI_Allreduce(MPI_IN_PLACE, &norm1, 1, MPI_DOUBLE, MPI_SUM, pmat.comm()); + MPI_Allreduce(MPI_IN_PLACE, &norm2, 1, MPI_DOUBLE, MPI_SUM, pmat.comm()); + norm1 = std::sqrt(norm1); + norm2 = std::sqrt(norm2); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "Frobenius norm"); + int local_flag = is_sym ? 1 : 0; + int global_flag = 0; + MPI_Allreduce(&local_flag, &global_flag, 1, MPI_INT, MPI_LAND, pmat.comm()); + is_sym = (global_flag != 0); +#else + assert(pmat.is_serial()); + std::vector minus_mat(pmat.get_local_size()); + std::vector sum_mat(pmat.get_local_size()); + for (std::size_t i = 0;i < n;++i) { + for (std::size_t j = i;j < n;++j) { + minus_mat[i * n + j] = mat[i * n + j] - mat[j * n + i]; + minus_mat[j * n + i] = -minus_mat[i * n + j]; + if (std::abs(minus_mat[i * n + j]) > threshold) {is_sym = false; } + sum_mat[i * n + j] = mat[i * n + j] + mat[j * n + i]; + sum_mat[j * n + i] = sum_mat[i * n + j]; + } + } + norm1 = LapackConnector::lange(norm_type, n, n, minus_mat.data(), n, &work_dummy); + norm2 = LapackConnector::lange(norm_type, n, n, sum_mat.data(), n, &work_dummy); +#endif + if (GlobalV::MY_RANK == 0) { + std::cout << "| Symmetric check: ||B - B^T||_F = " << norm1 << ", ||B + B^T||_F = " << norm2 << std::endl; + std::cout << "| ||B - B^T||_F / ||B + B^T||_F = " << norm1 / norm2 << std::endl; + } + return is_sym; } /// get the Psi wrapper of the selected spin from the Psi object @@ -142,37 +271,309 @@ namespace LR_Util psi_bfirst.fix_kb(ik_now, ib_now); return psi_kfirst; } +//=================2D-block Parallel=============== + /// @brief assign global X to 2d-matrix, its col is band(excition state), and row is { spin, k-point, occ, virt } + /// @attention pX is 2d-blocked as {occ, virt}, this assignment is used to calculate transition density matrix c_b X_{bj} c_j + /// @todo this function is a merge version of HamiltULR::global2local and HamiltLR::global2local, they should be replaced + template + void global2local_X(T* local_X, const T* global_X, const int nband, const int nk, + const std::vector& nocc, const std::vector& nvirt, const std::vector& pX, + const bool openshell) + { + const int nspin_X = openshell ? 2 : 1; + const std::vector npairs = { nocc[0] * nvirt[0], nocc[1] * nvirt[1] }; + const int gdim = openshell ? nk * (npairs[0] + npairs[1] ) : nk * npairs[0]; + const int ldim = openshell ? nk * (pX[0].get_local_size() + pX[1].get_local_size()) : nk * pX[0].get_local_size(); + + for (int ib = 0;ib < nband;++ib) + { + const int loffset_b = ib * ldim; + const int goffset_b = ib * gdim; + for (int is = 0;is < nspin_X;++is) + { + const int loffset_bs = loffset_b + is * nk * pX[0].get_local_size(); + const int goffset_bs = goffset_b + is * nk * npairs[0]; + for (int ik = 0;ik < nk;++ik) + { + const int loffset = loffset_bs + ik * pX[is].get_local_size(); + const int goffset = goffset_bs + ik * npairs[is]; + for (int lo = 0;lo < pX[is].get_col_size();++lo) + { + const int go = pX[is].local2global_col(lo); + for (int lv = 0;lv < pX[is].get_row_size();++lv) + { + const int gv = pX[is].local2global_row(lv); + local_X[loffset + lo * pX[is].get_row_size() + lv] = global_X[goffset + go * nvirt[is] + gv]; + } + } + } + } + } + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "global2local_X"); + } #ifdef __MPI + struct BlockHead + { + int ivirt; // {nvirt, nocc} global row + int iocc; // {nvirt, nocc} global col + int ik; + int src_rank; + void set(int iv_, int io_, int ik_, int src_rank_) + { + ivirt = iv_; + iocc = io_; + ik = ik_; + src_rank = src_rank_; + } + }; + inline MPI_Datatype mpi_type_blockhead() + { + static MPI_Datatype dt = MPI_DATATYPE_NULL; + static bool committed = false; + if (!committed) + { + int blen[4] = {1, 1, 1, 1}; + MPI_Aint disp[4]; + MPI_Datatype types[4] = {MPI_INT, MPI_INT, MPI_INT, MPI_INT}; + + disp[0] = static_cast(offsetof(BlockHead, ivirt)); + disp[1] = static_cast(offsetof(BlockHead, iocc)); + disp[2] = static_cast(offsetof(BlockHead, ik)); + disp[3] = static_cast(offsetof(BlockHead, src_rank)); + + MPI_Type_create_struct(4, blen, disp, types, &dt); + MPI_Type_commit(&dt); + committed = true; + } + return dt; + } + + /// @brief assign X in pA to X in pX + /// @note if use Cpxgemr2d, the communication time will be much more than MPIAlltoall by hand template - void gather_2d_to_full(const Parallel_2D& pv, const T* submat, T* fullmat, bool col_first, int global_nrow, int global_ncol) + void pA2pX(T* X_pX, const T* X_pA, const int nband, const int nk, + const std::vector& nocc, const std::vector& nvirt, + const std::vector& pX, const Parallel_2D& pA, + const int row_offset, const int col_offset, const bool openshell) { - ModuleBase::TITLE("LR_Util", "gather_2d_to_full"); - auto get_mpi_datatype = []() -> MPI_Datatype { - if (std::is_same::value) { return MPI_INT; } - if (std::is_same::value) { return MPI_FLOAT; } - else if (std::is_same::value) { return MPI_DOUBLE; } - if (std::is_same>::value) { return MPI_COMPLEX; } - else if (std::is_same>::value) { return MPI_DOUBLE_COMPLEX; } - else { throw std::runtime_error("gather_2d_to_full: unsupported type"); } - }; - - // zeros - for (int i = 0;i < global_nrow * global_ncol;++i) { fullmat[i] = 0.0; -} - //copy - for (int i = 0;i < pv.get_row_size();++i) { - for (int j = 0;j < pv.get_col_size();++j) { - if (col_first) { + ModuleBase::TITLE("LR_Util", "pA2pX"); + ModuleBase::timer::start("LR_Util", "pA2pX"); + const int nspin_X = openshell ? 2 : 1; + const std::vector npairs = { nocc[0] * nvirt[0], nocc[1] * nvirt[1] }; + const std::size_t gdim = openshell ? nk * (npairs[0] + npairs[1] ) : nk * npairs[0]; + const std::size_t ldim = openshell ? nk * (pX[0].get_local_size() + pX[1].get_local_size()) : nk * pX[0].get_local_size(); + assert(pA.get_global_row_size() == gdim || pA.get_global_row_size() == 2 * gdim); + assert(pA.get_global_col_size() == gdim || pA.get_global_col_size() == 2 * gdim); + for (int is = 0;is < nspin_X;++is) + { + assert(pX[is].get_global_row_size() == nvirt[is]); + assert(pX[is].get_global_col_size() == nocc[is]); + } + MPI_Datatype mpitype_blockhead = mpi_type_blockhead(); + + // 0. outer loop of row, communicate per 64 kpoints + constexpr int comm_nk = 64; + std::vector send_head_counts(GlobalV::NPROC, 0), recv_head_counts(GlobalV::NPROC, 0); + std::vector send_buffer_counts(GlobalV::NPROC, 0), recv_buffer_counts(GlobalV::NPROC, 0); + std::vector shdispls(GlobalV::NPROC, 0), rhdispls(GlobalV::NPROC, 0); // displacements of block heads + std::vector sbdispls(GlobalV::NPROC, 0), rbdispls(GlobalV::NPROC, 0); // displacements of buffer + std::vector cursor_head(GlobalV::NPROC, 0), cursor_buffer(GlobalV::NPROC, 0); + std::vector send_heads, recv_heads; + std::vector send_buffers, recv_buffers; + int max_send_head_total = 0, max_recv_head_total = 0; + int max_send_buffer_total = 0, max_recv_buffer_total = 0; + + // 1. pre-calculate local bands and global band list on each processor + const int gb_start = col_offset; + const int gb_end = col_offset + nband; + std::vector lnbands(GlobalV::NPROC, 0); + std::vector> gblist(GlobalV::NPROC); + const int nbpA = pA.get_block_size(); + for (int gb = gb_start; gb < gb_end; ++gb) + { + int proc_col = (gb / nbpA) % pA.dim1; + for (int pr = 0; pr < pA.dim0; ++pr) + { + int ownerA = proc_col + pr * pA.dim1; + ++lnbands[ownerA]; + gblist[ownerA].push_back(gb); + } + } + const int lnband = lnbands[GlobalV::MY_RANK]; + if (lnband == 0) + { + throw std::runtime_error("rank" + std::to_string(GlobalV::MY_RANK) + "lnband = 0, please check the pA and pX distribution!"); + } + const int lb_start = pA.global2local_col(gblist[GlobalV::MY_RANK].front()); + const int lb_end = lb_start + lnband; + + for (int is = 0; is < nspin_X; ++is) + { + const std::size_t gspin_base = is * nk * npairs[0]; + const std::size_t lspin_base = is * nk * pX[0].get_local_size(); + const std::size_t nrow = comm_nk * npairs[is]; + const std::size_t grow_start = row_offset + gspin_base; + const std::size_t grow_end = grow_start + nk * npairs[is]; + for (int irow_start = grow_start; irow_start < grow_end; irow_start += nrow) + { + std::fill(send_head_counts.begin(), send_head_counts.end(), 0); + std::fill(send_buffer_counts.begin(), send_buffer_counts.end(), 0); + + // 2. calculate and coummunicate block counts, then calculate block displacements + int irow_end = std::min(irow_start + nrow, grow_end); + for (int irow = irow_start; irow < irow_end; ++irow) + { + const int lrpA = pA.global2local_row(irow); + if (lrpA == -1) continue; + const int ir = irow - row_offset - gspin_base; + const int ik = ir / npairs[is]; + const int io = (ir-ik*npairs[is]) / nvirt[is]; + const int iv = ir % nvirt[is]; + int ownerX = pX[is].owner_processor(iv, io); + if (ownerX != GlobalV::MY_RANK) + { + ++send_head_counts[ownerX]; + } + } + for (int p = 0; p < GlobalV::NPROC; ++p) + { + std::size_t nbuffer = size_t(lnband) * send_head_counts[p]; + if (nbuffer > std::numeric_limits::max()) + { + throw std::overflow_error("in pA2pX: overflow converting to int!"); + } + send_buffer_counts[p] = static_cast(nbuffer); + } + + assert(send_head_counts.at(GlobalV::MY_RANK) == 0); + assert(send_buffer_counts.at(GlobalV::MY_RANK) == 0); + MPI_Alltoall(send_head_counts.data(), 1, MPI_INT, recv_head_counts.data(), 1, MPI_INT, pA.comm()); + MPI_Alltoall(send_buffer_counts.data(), 1, MPI_INT, + recv_buffer_counts.data(), 1, MPI_INT, pA.comm()); + + int send_head_total = 0, recv_head_total = 0, send_buffer_total = 0, recv_buffer_total = 0; + for (int p = 0; p < GlobalV::NPROC; ++p) + { + shdispls[p] = send_head_total; send_head_total += send_head_counts[p]; + rhdispls[p] = recv_head_total; recv_head_total += recv_head_counts[p]; + sbdispls[p] = send_buffer_total; send_buffer_total += send_buffer_counts[p]; + rbdispls[p] = recv_buffer_total; recv_buffer_total += recv_buffer_counts[p]; + } + if (send_head_total > max_send_head_total) + { max_send_head_total = send_head_total; send_heads.reserve(max_send_head_total); } + if (recv_head_total > max_recv_head_total) + { max_recv_head_total = recv_head_total; recv_heads.reserve(max_recv_head_total); } + if (send_buffer_total > max_send_buffer_total) + { max_send_buffer_total = send_buffer_total; send_buffers.reserve(max_send_buffer_total); } + if (recv_buffer_total > max_recv_buffer_total) + { max_recv_buffer_total = recv_buffer_total; recv_buffers.reserve(max_recv_buffer_total); } + + // 3. prepare block heads and buffers + send_heads.resize(send_head_total); + recv_heads.resize(recv_head_total); + send_buffers.resize(send_buffer_total); + recv_buffers.resize(recv_buffer_total); + cursor_head = shdispls; + cursor_buffer = sbdispls; + + for (int irow = irow_start; irow < irow_end; ++irow) + { + const int lr = pA.global2local_row(irow); + if (lr == -1) continue; + const std::size_t ir = irow - row_offset - gspin_base; + const std::size_t ik = ir / npairs[is]; + const int io = (ir-ik*npairs[is]) / nvirt[is]; + const int iv = ir % nvirt[is]; + const int ownerX = pX[is].owner_processor(iv, io); + if (ownerX == GlobalV::MY_RANK) + { + for (std::size_t lb = lb_start; lb < lb_end; ++lb) + { + const int gb = pA.local2global_col(lb); + const std::size_t lb_base = std::size_t(gb - col_offset) * std::size_t(ldim); + const std::size_t lrowX = pX[is].global2local_row(iv); + const std::size_t lcolX = pX[is].global2local_col(io); + const std::size_t lX_indx = lrowX + lcolX * pX[is].get_row_size() + lb_base + + lspin_base + ik * pX[is].get_local_size(); + X_pX[lX_indx] = X_pA[static_cast(lr) + lb * pA.get_row_size()]; + } + } + else + { + BlockHead& head = send_heads[cursor_head[ownerX]++]; + head.set(iv, io, ik, GlobalV::MY_RANK); + for (std::size_t lb = lb_start; lb < lb_end; ++lb) + { + send_buffers[cursor_buffer[ownerX]++] = X_pA[static_cast(lr) + lb * pA.get_row_size()]; + } + } + } + for (int p = 0; p < GlobalV::NPROC; ++p) + { + assert(cursor_head[p] == shdispls[p] + send_head_counts[p]); + assert(cursor_buffer[p] == sbdispls[p] + send_buffer_counts[p]); + } + // 4. communicate block heads and buffers + MPI_Alltoallv(send_heads.data(), send_head_counts.data(), shdispls.data(), mpitype_blockhead, + recv_heads.data(), recv_head_counts.data(), rhdispls.data(), mpitype_blockhead, + pA.comm()); + MPI_Alltoallv(send_buffers.data(), send_buffer_counts.data(), sbdispls.data(), LR_Util::MPIType::value(), + recv_buffers.data(), recv_buffer_counts.data(), rbdispls.data(), LR_Util::MPIType::value(), + pA.comm()); + + // 5. unpack received buffers to pX + for (int src = 0; src < GlobalV::NPROC; ++src) + { + const int nh = recv_head_counts[src]; + if (nh == 0) continue; + int buf_cursor = rbdispls[src]; + const int ib_start = rhdispls[src]; + const int ib_end = ib_start + nh; + for (int iblock = ib_start; iblock < ib_end; ++iblock) + { + const BlockHead& rh = recv_heads[iblock]; + assert(rh.src_rank == src); + const std::size_t ik = rh.ik; + const std::size_t lr = pX[is].global2local_row(rh.ivirt); + const std::size_t lc = pX[is].global2local_col(rh.iocc); + for (int ib = 0; ib < lnbands[src]; ++ib) + { + const std::size_t lb_base = std::size_t(gblist[src][ib] - col_offset) * std::size_t(ldim); + const std::size_t lX_indx = lr + lc * pX[is].get_row_size() + lb_base + + lspin_base + ik * pX[is].get_local_size(); + X_pX[lX_indx] = recv_buffers[ib + buf_cursor]; + } + buf_cursor += lnbands[src]; + } + assert(buf_cursor == rbdispls[src] + recv_buffer_counts[src]); + } + } // end of irow_start + } // end of is + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "pA2pX"); + ModuleBase::timer::end("LR_Util", "pA2pX"); + } + + template + void gather_2d_to_full(const Parallel_2D& pv, const T* submat, T* fullmat, + const bool row_major, const std::size_t global_nrow, const std::size_t global_ncol) + { + assert(pv.get_global_row_size() == global_nrow); + assert(pv.get_global_col_size() == global_ncol); + // copy +#ifdef _OPENMP +#pragma omp parallel for collapse(2) +#endif + for (int j = 0;j < pv.get_col_size();++j) { + for (int i = 0;i < pv.get_row_size();++i) { + if (row_major) { fullmat[pv.local2global_row(i) * global_ncol + pv.local2global_col(j)] = submat[i * pv.get_col_size() + j]; } else { fullmat[pv.local2global_col(j) * global_nrow + pv.local2global_row(i)] = submat[j * pv.get_row_size() + i]; -} -} -} - - //reduce to root - MPI_Allreduce(MPI_IN_PLACE, fullmat, global_nrow * global_ncol, get_mpi_datatype(), MPI_SUM, pv.comm()); + } + } + } + MPI_Allreduce(MPI_IN_PLACE, fullmat, global_nrow * global_ncol, LR_Util::MPIType::value(), MPI_SUM, pv.comm()); }; #endif diff --git a/source/source_lcao/module_lr/utils/lr_util_hcontainer.h b/source/source_lcao/module_lr/utils/lr_util_hcontainer.h index 6b4ee697e5..e75b13ff43 100644 --- a/source/source_lcao/module_lr/utils/lr_util_hcontainer.h +++ b/source/source_lcao/module_lr/utils/lr_util_hcontainer.h @@ -107,13 +107,14 @@ namespace LR_Util for (int iR = 0;iR < ap1->get_R_size();++iR) { const ModuleBase::Vector3& R = ap1->get_R_index(iR); + // std::cout<< "dot_R_matrix: iat1=" << iat1 << ", iat2=" << iat2 << ", R=(" << R.x << ", " << R.y << ", " << R.z << ")"<get_HR_values(R.x, R.y, R.z); auto mat2 = ap2->get_HR_values(R.x, R.y, R.z); - sum += std::inner_product(mat1.get_pointer(), mat1.get_pointer() + mat1.get_memory_size(), mat2.get_pointer(), (TR1)0.0); + sum += std::inner_product(mat1.get_pointer(), mat1.get_pointer() + mat1.get_col_size()*mat1.get_row_size(), mat2.get_pointer(), (TR1)0.0); } } } - Parallel_Reduce::reduce_all(sum); + // Parallel_Reduce::reduce_all(sum); // not needed, since it will be reduced outside return sum; } diff --git a/source/source_lcao/module_lr/utils/lr_util_print.h b/source/source_lcao/module_lr/utils/lr_util_print.h index b269e76e01..789306b9bd 100644 --- a/source/source_lcao/module_lr/utils/lr_util_print.h +++ b/source/source_lcao/module_lr/utils/lr_util_print.h @@ -25,6 +25,10 @@ namespace LR_Util int read_value(const std::string& file, T* ptr, const int& size, Args&&... args) { std::ifstream ifs(file); + if (!ifs.is_open()){ + std::cerr << "Error: Cannot open file " << file << std::endl; + return -1; + } const int res = read_value(ifs, ptr, size, args...); ifs.close(); return res; @@ -49,6 +53,10 @@ namespace LR_Util int write_value(const std::string& file, const int& prec, const T* ptr, const int& size, Args&&... args) { std::ofstream ofs(file); + if (!ofs.is_open()){ + std::cerr << "Error: Cannot open file " << file << std::endl; + return -1; + } ofs << std::setprecision(prec) << std::scientific; const int res = write_value(ofs, ptr, size, args...); ofs.close(); diff --git a/source/source_lcao/module_lr/utils/mo_type.h b/source/source_lcao/module_lr/utils/mo_type.h new file mode 100644 index 0000000000..5f38f527e8 --- /dev/null +++ b/source/source_lcao/module_lr/utils/mo_type.h @@ -0,0 +1,43 @@ +#pragma once +#include +namespace LR_Util +{ +#ifndef MO_TYPE_H +#define MO_TYPE_H + enum MO_TYPE { OO, VO, OV, VV, ALL }; +#endif +/* the first index is contiguous in memory +MO_TYPE: OO VO OV VV ALL +nmo1 nocc nocc nvirt nvirt nocc+nvirt +nmo2 nocc nvirt nocc nvirt nocc+nvirt +imo1 0 0 nocc nocc 0 +imo2 0 nocc 0 nocc 0 +*/ +inline void set_dim(const MO_TYPE type, const int& nocc, const int& nvirt, + int& nmo1, int& nmo2, int& imo1, int& imo2) +{ + switch(type) + { + case MO_TYPE::OO: + nmo1 = nocc; nmo2 = nocc; imo1 = 0; imo2 = 0; + break; + case MO_TYPE::VO: + nmo1 = nocc; nmo2 = nvirt; imo1 = 0; imo2 = nocc; + break; + case MO_TYPE::OV: + nmo1 = nvirt; nmo2 = nocc; imo1 = nocc; imo2 = 0; + break; + case MO_TYPE::VV: + nmo1 = nvirt; nmo2 = nvirt; imo1 = nocc; imo2 = nocc; + break; + case MO_TYPE::ALL: + nmo1 = nocc + nvirt; + nmo2 = nocc + nvirt; + imo1 = 0; + imo2 = 0; + break; + default: + throw std::runtime_error("Error in LR::set_dim: unknown MO_TYPE"); + } +} +} \ No newline at end of file diff --git a/source/source_lcao/module_lr/utils/spectrum_mo.hpp b/source/source_lcao/module_lr/utils/spectrum_mo.hpp new file mode 100644 index 0000000000..d0b3e37275 --- /dev/null +++ b/source/source_lcao/module_lr/utils/spectrum_mo.hpp @@ -0,0 +1,208 @@ +#pragma once +#include "source_base/tool_title.h" +#include "source_basis/module_nao/two_center_bundle.h" +#include "source_cell/klist.h" +#include "source_io/module_parameter/parameter.h" +#include "source_lcao/module_hcontainer/hcontainer_funcs.h" +#include "source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h" +#include "source_lcao/module_lr/utils/lr_util.h" +#include "source_lcao/module_lr/utils/lr_util_print.h" +#include "source_lcao/module_rt/velocity_op.h" +#include "source_psi/psi.h" + +namespace LR_Util +{ +template +std::vector> cal_velocity_mo(const UnitCell& ucell, + const Grid_Driver& gd, + const TwoCenterBundle& two_center_bundle, + const Parallel_Orbitals& pmat,/*nbasis×nbasis*/ + const Parallel_2D& pc,/*nbasis×nbands*/ + const K_Vectors& kv, + const psi::Psi& psi_ks, + const int nk, + const int nspin_tmp, + const int nbasis, + const std::vector nocc, + const std::vector nvirt) +{ + ModuleBase::TITLE("LR_Util", "cal_velocity_mo"); + ModuleBase::timer::start("LR_Util", "cal_velocity_mo"); + std::cout<<"Calculating velocity matrix in KS presentation..."<> vR(&ucell, &gd, &pmat, orb, two_center_bundle.overlap_orb.get()); + vR.calculate_grad_term(); // $<\mu, 0|-i∇r|\nu, R>$ + vR.calculate_vcomm_r(); // $<\mu, 0|i[Vnl, r]|\nu, R>$ + + int nks = kv.get_nks(); // include spin + assert(nks == nk * nspin_tmp); + assert(psi_ks.get_nk() == nks); + int KS_num = nocc[0] + nvirt[0]; + std::vector> velocity_mo(nspin_tmp * 3 * nk * KS_num * KS_num, 0.0); + Parallel_2D pmo; + LR_Util::setup_2d_division(pmo, pmat.get_block_size(), KS_num, KS_num +#ifdef __MPI + , pc.blacs_ctxt +#endif + ); + + //1. psi_ks to c_psi_ks>, ensure complex for dipole calculation + psi::Psi> c_psi_ks(nks, + pc.get_col_size(), + pc.get_row_size(), + kv.ngk, + true); + for(int iks = 0; iks < nks; ++iks) + { + for(int ic = 0; ic < pc.get_col_size(); ++ic)//band + { + for(int ir = 0; ir < pc.get_row_size(); ++ir)//basis + { + c_psi_ks(iks, ic, ir) = std::complex(psi_ks(iks, ic, ir)); + } + } + } + + //2. calculate v_mo = c^\dagger v c + std::vector vk(nks, LR_Util::newTensor>({ pmat.get_col_size(), pmat.get_row_size() })); + for (int id = 0; id < 3; ++id) + { + for (auto& v : vk) v.zero(); + + std::vector> v_mo(nks * pmo.get_local_size(), 0.0); + + for (int is = 0; is < nspin_tmp; ++is) + { + assert(KS_num == nocc[is] + nvirt[is]); + + for (int ik = is*nk; ik < (is+1)*nk; ++ik) + { + hamilt::folding_HR(*vR.get_current_term_pointer(id), vk[ik].data>(), kv.kvec_d[ik], pmat.get_row_size(), 1/*column-major*/); + } + } +#ifdef __MPI + LR::ao_to_mo_pblas(vk, pmat, c_psi_ks, pc, nbasis, + nocc[0], nvirt[0], pmo, v_mo.data(), + false, // add_on + LR_Util::MO_TYPE::ALL); +#else + LR::ao_to_mo_blas(vk, c_psi_ks, + nocc[0], nvirt[0], v_mo.data(), + false , //add_on + LR_Util::MO_TYPE::ALL); +#endif + // gather local vk to global velocity_mo + for (int is = 0; is < nspin_tmp; ++is) + { + for (int ik = 0; ik < nk; ++ik) + { + int glb_offset = (is * 3 * nk + id * nk + ik) * KS_num * KS_num; + int loc_offset = ik * pmo.get_local_size(); + for (int j = 0; j < pmo.get_col_size(); ++j){ + for (int i = 0; i < pmo.get_row_size(); ++i){ + velocity_mo[glb_offset + pmo.local2global_col(j) * KS_num + pmo.local2global_row(i)] + = v_mo[loc_offset + j * pmo.get_row_size() + i]; + } + } + + //std::cout<< "is" << is << "id: " << id << " ik: " << ik << " v_mo: " << std::endl; + //LR_Util::print_value(velocity_mo.data()+(is * 3 * nk + id * nk + ik) * KS_num * KS_num, KS_num, KS_num); + } + } + }//id + MPI_Allreduce(MPI_IN_PLACE, velocity_mo.data(), velocity_mo.size(), LR_Util::MPIType::value(), MPI_SUM, pmo.comm()); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "Finish velocity matrix in KS presentation."); + ModuleBase::timer::end("LR_Util", "cal_velocity_mo"); + return velocity_mo; +} + +/// @brief output the velocity matrix in KS presentation in human-read friendly format +inline void output_spectrum_mo(const std::vector>& out_spectrum_mo, + const std::string& filename, + const double* const eig_ks, + const int nk, + const int nspin_tmp, + const int KS_num, + const K_Vectors& kv) +{ + assert(out_spectrum_mo.size() == nspin_tmp * 3 * nk * KS_num * KS_num); + std::ofstream ofs(filename); + ofs << "Data Unit (Hartree * Bohr)" << std::endl; + ofs << "NOTICE: KS_index are restricted in nocc and nvirt" << std::endl; + + int step = nk * KS_num * KS_num; + int offset = 0, ipair = 0; + for (int is = 0; is < nspin_tmp; ++is) + { + ofs << "ispin: " << is << std::endl; + for (int ik = 0;ik < nk;++ik) + { + ofs << "k-point: " << ik << " " << kv.kvec_d[ik] << std::endl; + ofs << std::setw(4) << "KS1" << std::setw(10) << "E1(eV)" << std::setw(6) << "KS2" << std::setw(10) << "E2(eV)" + << std::setw(16) << "x" << std::setw(23) << "|x|^2" << std::setw(18) << "y" << std::setw(23) <<"|y|^2" + << std::setw(18) << "z" << std::setw(23) <<"|z|^2" << std::setw(13) << "average" << std::endl; + + offset = (is * 3 * nk + ik) * KS_num * KS_num; + for (int i = 0; i < KS_num; ++i) + { + for (int j = i; j < KS_num ; ++j) + { + ipair = offset + i * KS_num + j; + double average = (std::norm(out_spectrum_mo[ipair]) + std::norm(out_spectrum_mo[ipair + step]) + std::norm(out_spectrum_mo[ipair + 2 * step])) / 3.0; + ofs << std::setw(3) << i << std::setw(12) << std::setprecision(6) << eig_ks[i + ik*KS_num] * ModuleBase::Ry_to_eV + << std::setw(4) << j << std::setw(12) << eig_ks[j + ik*KS_num] * ModuleBase::Ry_to_eV + << std::setw(28) << out_spectrum_mo[ipair] << std::setw(13) << std::norm(out_spectrum_mo[ipair]) + << std::setw(28) << out_spectrum_mo[ipair + step] << std::setw(13) << std::norm(out_spectrum_mo[ipair + step]) + << std::setw(28) << out_spectrum_mo[ipair + 2 * step] << std::setw(13) << std::norm(out_spectrum_mo[ipair + 2 * step] ) + << std::setw(13) << average << std::endl; + } + } + } + } + ofs.close(); +} + +/// @brief output the velocity matrix in KS presentation in LibRPA format +inline void output_spectrum_mo_librpa(const std::vector>& out_spectrum_mo, + const std::string& filename, + const int nk, + const int nspin_tmp, + const int KS_num, + const K_Vectors& kv) +{ + assert(out_spectrum_mo.size() == nspin_tmp * 3 * nk * KS_num * KS_num); + std::ofstream ofs(filename); + ofs << std::scientific << nk << std::endl; + ofs << nspin_tmp << std::endl; + ofs << PARAM.inp.nbands << std::endl; + ofs << PARAM.globalv.nlocal << std::endl; + double HaBohrToEvAng = ModuleBase::Hartree_to_eV * ModuleBase::BOHR_TO_A; // 27.211396 * 0.5291770 + int offset = 0, ipair = 0; + for (int is = 0; is < nspin_tmp; ++is) + { + for (int ik = 0; ik < nk; ++ik) + { + for (int id = 0; id < 3; ++id) + { + offset = (is * 3 * nk + id * nk + ik) * KS_num * KS_num; + ofs <<" " << id + 1 << " " << ik + 1 << " " << is + 1 << std::endl; + for (int ib1 = 0; ib1 < KS_num; ++ib1) + { + for (int ib2 = 0; ib2 < KS_num; ++ib2) + { + ipair = offset + ib1 * KS_num + ib2; + ofs << std::setw(26) << std::fixed << std::setprecision(16) << out_spectrum_mo[ipair].real() * HaBohrToEvAng + << std::setw(26) << out_spectrum_mo[ipair].imag() * HaBohrToEvAng << std::endl; + } + } + } + + } + } +} +} \ No newline at end of file diff --git a/source/source_lcao/module_ri/LRI_CV_Tools.h b/source/source_lcao/module_ri/LRI_CV_Tools.h index 55f01db3ec..52db749376 100644 --- a/source/source_lcao/module_ri/LRI_CV_Tools.h +++ b/source/source_lcao/module_ri/LRI_CV_Tools.h @@ -143,6 +143,8 @@ using TLRI = std::map>>; template TLRI read_Cs_ao(const std::string& file_path, const double& threshold = 1e-10); template +TLRI read_Cs_ao_all(const std::string& path, const double& threshold = 1e-10); +template void write_Cs_ao(const TLRI& Vs, const std::string& file_path); template TLRI read_Vs_abf(const std::string& file_path, const double& threshold = 1e-10); diff --git a/source/source_lcao/module_ri/RPA_LRI.h b/source/source_lcao/module_ri/RPA_LRI.h index 36b4c5fd57..b67bc90e49 100644 --- a/source/source_lcao/module_ri/RPA_LRI.h +++ b/source/source_lcao/module_ri/RPA_LRI.h @@ -68,7 +68,12 @@ template class RPA_LRI void out_eigen_vector(const Parallel_Orbitals& parav, const psi::Psi& psi); void out_struc(const UnitCell& ucell); void out_bands(const elecstate::ElecState *pelec); - + void out_velocity(const UnitCell& ucell, + const Grid_Driver& gd, + const TwoCenterBundle& two_center_bundle, + const Parallel_Orbitals& parav,/*nbasis×nbasis*/ + const psi::Psi& psi, + const elecstate::ElecState* pelec); void output_cut_coulomb_cs(const UnitCell& ucell, Exx_LRI* exx_lri_rpa); void out_Cs(const UnitCell& ucell, std::map>>& Cs_in, std::string filename); void out_coulomb_k(const UnitCell& ucell, @@ -83,6 +88,7 @@ template class RPA_LRI Tdata Erpa; private: + const std::string& outdir = PARAM.inp.rpa_outdir; Exx_Info_RI info; const K_Vectors *p_kv=nullptr; MPI_Comm mpi_comm; diff --git a/source/source_lcao/module_ri/RPA_LRI.hpp b/source/source_lcao/module_ri/RPA_LRI.hpp index 2817390201..309912c174 100644 --- a/source/source_lcao/module_ri/RPA_LRI.hpp +++ b/source/source_lcao/module_ri/RPA_LRI.hpp @@ -15,9 +15,11 @@ #include "source_lcao/module_ri/module_exx_symmetry/symmetry_rotation.h" #include "RPA_LRI.h" +#include "source_base/global_function.h" #include "source_basis/module_ao/element_basis_index-ORB.h" #include "source_estate/elecstate_lcao.h" #include "source_io/module_parameter/parameter.h" +#include "source_lcao/module_lr/utils/spectrum_mo.hpp" #if defined(__GLIBC__) #include @@ -45,6 +47,7 @@ void RPA_LRI::postSCF(const UnitCell& ucell, { ModuleBase::TITLE("RPA_LRI", "postSCF"); ModuleBase::timer::start("RPA_LRI", "postSCF"); + ModuleBase::GlobalFunc::MAKE_DIR(outdir); this->cal_postSCF_exx(dm, mpi_comm_in, ucell, kv, orb); this->init(mpi_comm_in, kv, orb.cutoffs()); @@ -624,7 +627,7 @@ void RPA_LRI::out_abfs_overlap(const UnitCell& ucell, ss << filename << GlobalV::MY_RANK << ".txt"; std::ofstream ofs; - ofs.open(ss.str().c_str(), std::ios::out); + ofs.open(outdir + ss.str().c_str(), std::ios::out); ofs << nks_tot << std::endl; @@ -972,8 +975,9 @@ void RPA_LRI::out_eigen_vector(const Parallel_Orbitals& parav, const p std::ofstream ofs; if (GlobalV::MY_RANK == 0) { - ofs.open(ss.str().c_str(), std::ios::out); + ofs.open(outdir + ss.str().c_str(), std::ios::out); } + ofs << std::fixed << std::setprecision(15); std::vector is_wfc_ib_iw(npsin_tmp); for (int is = 0; is < npsin_tmp; is++) { @@ -1009,8 +1013,8 @@ void RPA_LRI::out_eigen_vector(const Parallel_Orbitals& parav, const p { for (int is = 0; is < npsin_tmp; is++) { - ofs << std::setw(30) << std::fixed << std::setprecision(15) << is_wfc_ib_iw[is](ib, iw).real() - << std::setw(30) << std::fixed << std::setprecision(15) << is_wfc_ib_iw[is](ib, iw).imag() + ofs << std::setw(30) << is_wfc_ib_iw[is](ib, iw).real() + << std::setw(30) << is_wfc_ib_iw[is](ib, iw).imag() << std::endl; } } @@ -1035,7 +1039,8 @@ void RPA_LRI::out_struc(const UnitCell& ucell) std::stringstream ss; ss << "stru_out"; std::ofstream ofs; - ofs.open(ss.str().c_str(), std::ios::out); + ofs.open(outdir + ss.str().c_str(), std::ios::out); + ofs << std::fixed << std::setprecision(9); ofs << lat.e11 << std::setw(15) << lat.e12 << std::setw(15) << lat.e13 << std::endl; ofs << lat.e21 << std::setw(15) << lat.e22 << std::setw(15) << lat.e23 << std::endl; ofs << lat.e31 << std::setw(15) << lat.e32 << std::setw(15) << lat.e33 << std::endl; @@ -1059,8 +1064,7 @@ void RPA_LRI::out_struc(const UnitCell& ucell) : ucell.atoms[it].tau[ia].y; const double& z = direct ? ucell.atoms[it].tau[ia].z * ucell.lat0 : ucell.atoms[it].tau[ia].z; - ofs << std::setw(15) << std::fixed << std::setprecision(9) << x << std::setw(15) << std::fixed - << std::setprecision(9) << y << std::setw(15) << std::fixed << std::setprecision(9) << z + ofs << std::setw(15) << x << std::setw(15) << y << std::setw(15) << z << std::setw(15) << 1 << std::endl; } } @@ -1069,9 +1073,9 @@ void RPA_LRI::out_struc(const UnitCell& ucell) for (int ik = 0; ik != nks_tot; ik++) { - ofs << std::setw(15) << std::fixed << std::setprecision(9) << p_kv->kvec_c[ik].x * TWOPI_Bohr2A << std::setw(15) - << std::fixed << std::setprecision(9) << p_kv->kvec_c[ik].y * TWOPI_Bohr2A << std::setw(15) << std::fixed - << std::setprecision(9) << p_kv->kvec_c[ik].z * TWOPI_Bohr2A << std::endl; + ofs << std::setw(15) << p_kv->kvec_c[ik].x * TWOPI_Bohr2A << std::setw(15) + << p_kv->kvec_c[ik].y * TWOPI_Bohr2A << std::setw(15) + << p_kv->kvec_c[ik].z * TWOPI_Bohr2A << std::endl; } // added for BZ to IBZ (actually LibRPA interface only support BZ by 2025/03/30) if (PARAM.inp.symmetry == "-1") @@ -1098,7 +1102,8 @@ void RPA_LRI::out_bands(const elecstate::ElecState* pelec) std::stringstream ss; ss << "band_out"; std::ofstream ofs; - ofs.open(ss.str().c_str(), std::ios::out); + ofs.open(outdir + ss.str().c_str(), std::ios::out); + ofs << std::fixed << std::setprecision(15); ofs << nks_tot << std::endl; ofs << nspin_tmp << std::endl; ofs << PARAM.inp.nbands << std::endl; @@ -1113,8 +1118,8 @@ void RPA_LRI::out_bands(const elecstate::ElecState* pelec) for (int ib = 0; ib != PARAM.inp.nbands; ib++) { ofs << std::setw(5) << ib + 1 << " " << std::setw(8) << pelec->wg(ik + is * nks_tot, ib) * nks_tot - << std::setw(25) << std::fixed << std::setprecision(15) << pelec->ekb(ik + is * nks_tot, ib) / 2.0 - << std::setw(25) << std::fixed << std::setprecision(15) + << std::setw(25) << pelec->ekb(ik + is * nks_tot, ib) / 2.0 + << std::setw(25) << pelec->ekb(ik + is * nks_tot, ib) * ModuleBase::Ry_to_eV << std::endl; } } @@ -1132,8 +1137,9 @@ void RPA_LRI::out_Cs(const UnitCell& ucell, std::map::out_Cs(const UnitCell& ucell, std::map::out_coulomb_k(const UnitCell& ucell, ss << filename << GlobalV::MY_RANK << ".txt"; std::ofstream ofs; - ofs.open(ss.str().c_str(), std::ios::out); + ofs.open(outdir + ss.str().c_str(), std::ios::out); ofs << nks_tot << std::endl; + ofs << std::fixed << std::setprecision(15); for (auto& Ip: Vs) { auto I = Ip.first; @@ -1233,8 +1240,8 @@ void RPA_LRI::out_coulomb_k(const UnitCell& ucell, ofs << ik + 1 << " " << p_kv->wk[ik] / 2.0 * PARAM.inp.nspin << std::endl; for (int i = 0; i != vq_J.data->size(); i++) { - ofs << std::setw(25) << std::fixed << std::setprecision(15) << (*vq_J.data)[i].real() - << std::setw(25) << std::fixed << std::setprecision(15) << (*vq_J.data)[i].imag() << std::endl; + ofs << std::setw(25) << (*vq_J.data)[i].real() + << std::setw(25) << (*vq_J.data)[i].imag() << std::endl; } } } @@ -1243,6 +1250,43 @@ void RPA_LRI::out_coulomb_k(const UnitCell& ucell, ModuleBase::timer::end("RPA_LRI", "out_coulomb_k"); } +template +void RPA_LRI::out_velocity(const UnitCell &ucell, + const Grid_Driver &gd, + const TwoCenterBundle &two_center_bundle, + const Parallel_Orbitals ¶v,/*nbasis×nbasis*/ + const psi::Psi &psi, + const elecstate::ElecState* pelec) +{ + ModuleBase::TITLE("DFT_RPA_interface", "out_velocity"); + ModuleBase::timer::start("RPA_LRI", "out_velocity"); + + Parallel_2D parac;/*nbasis×nbands*/ + LR_Util::setup_2d_division(parac, parav.get_block_size(), PARAM.globalv.nlocal, PARAM.inp.nbands + #ifdef __MPI + , parav.blacs_ctxt + #endif + ); + + const int nk = PARAM.inp.nspin == 2 ? p_kv->get_nks() / 2 : p_kv->get_nks(); + const int nspin_tmp = PARAM.inp.nspin == 2 ? 2 : 1; + + // nocc and nvirt dosen't matter, their sum is actually used + std::vector nocc(2, PARAM.inp.nbands); + std::vector nvirt(2, 0); + + std::vector> velocity_mo = LR_Util::cal_velocity_mo(ucell, gd, two_center_bundle, + parav, parac, *this->p_kv, psi, nk, nspin_tmp, PARAM.globalv.nlocal, nocc, nvirt); + if (GlobalV::MY_RANK == 0){ + // for librpa readable + LR_Util::output_spectrum_mo_librpa(velocity_mo, outdir + "velocity_matrix", + nk, nspin_tmp, PARAM.inp.nbands, *this->p_kv); + // for human readable + // LR_Util::output_spectrum_mo(velocity_mo, PARAM.globalv.global_out_dir + "velocity_matrix_rpa.dat", + // pelec->ekb.c, nk, nspin_tmp, PARAM.inp.nbands, *this->p_kv); + } + ModuleBase::timer::end("RPA_LRI", "out_velocity"); +} // template // void RPA_LRI::init(const MPI_Comm &mpi_comm_in) diff --git a/source/source_lcao/module_ri/write_ri_cv.hpp b/source/source_lcao/module_ri/write_ri_cv.hpp index a90fa3c490..0a411fba11 100644 --- a/source/source_lcao/module_ri/write_ri_cv.hpp +++ b/source/source_lcao/module_ri/write_ri_cv.hpp @@ -1,6 +1,9 @@ // #include "source_lcao/module_ri/LRI_CV_Tools.h" +#include "source_base/global_function.h" +#include "source_base/tool_title.h" #include #include +#include #define IZ(x) int x = 0; #define SZ(x) std::size_t x = 0; @@ -38,7 +41,13 @@ namespace LRI_CV_Tools infile >> ia1 >> ia2 >> ic_1 >> ic_2 >> ic_3 >> nw1 >> nw2 >> nabf; const TC& box = { ic_1, ic_2, ic_3 }; RI::Tensor tensor_cs({ nabf, nw1, nw2 }); - for (std::size_t i = 0; i != nw1; i++) { for (std::size_t j = 0; j != nw2; j++) { for (std::size_t mu = 0; mu != nabf; mu++) { infile >> tensor_cs(mu, i, j); } } } + for (std::size_t i = 0; i != nw1; i++) { + for (std::size_t j = 0; j != nw2; j++) { + for (std::size_t mu = 0; mu != nabf; mu++) { + infile >> tensor_cs(mu, i, j); + } + } + } // no screening for data-structure consistency if (absmax(tensor_cs) >= threshold) { Cs[ia1 - 1][{ia2 - 1, box}] = tensor_cs; } } @@ -51,7 +60,13 @@ namespace LRI_CV_Tools { std::ofstream outfile; outfile.open(file_path); - outfile << Cs.size() << " " << Cs.at(0).size() / Cs.size() << std::endl; //natom, ncell + if (Cs.size() == 0) + { + outfile << " Empty Cs " << std::endl; + outfile.close(); + return; + } + outfile << Cs.size() << " " << Cs.begin()->second.size() / Cs.size() << std::endl; //natom, (ncell, not real if Cs is incomplete) for (auto& it1 : Cs) { const int& ia1 = it1.first; @@ -76,6 +91,47 @@ namespace LRI_CV_Tools outfile.close(); } + template + TLRI read_Cs_ao_all(const std::string& path, const double& threshold) + { + IZ(natom) IZ(ncell) IZ(ia1) IZ(ia2) IZ(ic_1) IZ(ic_2) IZ(ic_3) + SZ(nw1) SZ(nw2) SZ(nabf) + TLRI Cs; + + struct dirent *ptr; + DIR *dir; + dir = opendir(path.c_str()); + while ((ptr = readdir(dir)) != NULL){// read all the files in the directory + std::string fm(ptr->d_name); + if (fm.find("Cs_data_") == 0)// find file Cs_data_xxx + { + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "found Cs file: " + fm + ", start reading..."); + std::ifstream infile(path + fm); + infile >> natom >> ncell; // no use of ncell + + while (infile.peek() != EOF) + { + if(!(infile >> ia1 >> ia2 >> ic_1 >> ic_2 >> ic_3 >> nw1 >> nw2 >> nabf)){break;} + const TC& box = { ic_1, ic_2, ic_3 }; + RI::Tensor tensor_cs({ nabf, nw1, nw2 }); + for (std::size_t i = 0; i != nw1; i++) { + for (std::size_t j = 0; j != nw2; j++) { + for (std::size_t mu = 0; mu != nabf; mu++) { + infile >> tensor_cs(mu, i, j); + } + } + } + // no screening for data-structure consistency + if (absmax(tensor_cs) >= threshold) { Cs[ia1 - 1][{ia2 - 1, box}] = tensor_cs; } + } + infile.close(); + } + } + closedir(dir); + ModuleBase::TITLE("LRI_CV_Tools", "read_Cs_ao_all done."); + return Cs; + } + template TLRI read_Vs_abf(const std::string& file_path, const double& threshold) { @@ -106,7 +162,13 @@ namespace LRI_CV_Tools { std::ofstream outfile; outfile.open(file_path); - outfile << Vs.size() << " " << Vs.at(0).size() / Vs.size() << std::endl; //natom, ncell + if (Vs.size() == 0) + { + outfile << " Empty Vs " << std::endl; + outfile.close(); + return; + } + outfile << Vs.size() << " " << Vs.begin()->second.size() / Vs.size() << std::endl; //natom, (ncell, not real if Vs is incomplete) for (const auto& it1 : Vs) { const int& ia1 = it1.first; diff --git a/tests/08_EXX/51_GO_LR/KPT b/tests/08_EXX/51_GO_LR/KPT new file mode 100644 index 0000000000..c289c0158a --- /dev/null +++ b/tests/08_EXX/51_GO_LR/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 1 0 0 0 diff --git a/tests/08_EXX/51_GO_LR/result.ref b/tests/08_EXX/51_GO_LR/result.ref index 5c65fd6538..38e23a0cc9 100644 --- a/tests/08_EXX/51_GO_LR/result.ref +++ b/tests/08_EXX/51_GO_LR/result.ref @@ -1 +1,5 @@ -totexcitationenergyref 2.510666 \ No newline at end of file +excitationenergyref1 0.587373 +excitationenergyref2 0.727934 +excitationenergyref3 0.531918 +excitationenergyref4 0.663441 +totaltimeref 1.74 diff --git a/tests/08_EXX/52_GO_LR_PBE/KPT b/tests/08_EXX/52_GO_LR_PBE/KPT new file mode 100644 index 0000000000..c289c0158a --- /dev/null +++ b/tests/08_EXX/52_GO_LR_PBE/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 1 0 0 0 diff --git a/tests/08_EXX/52_GO_LR_PBE/result.ref b/tests/08_EXX/52_GO_LR_PBE/result.ref index 0086694b94..a6e3b42a47 100644 --- a/tests/08_EXX/52_GO_LR_PBE/result.ref +++ b/tests/08_EXX/52_GO_LR_PBE/result.ref @@ -1 +1,5 @@ -totexcitationenergyref 2.504802 +excitationenergyref1 0.589637 +excitationenergyref2 0.731349 +excitationenergyref3 0.526037 +excitationenergyref4 0.657779 +totaltimeref 1.82 diff --git a/tests/08_EXX/53_GO_LR_HF/KPT b/tests/08_EXX/53_GO_LR_HF/KPT new file mode 100644 index 0000000000..c289c0158a --- /dev/null +++ b/tests/08_EXX/53_GO_LR_HF/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 1 0 0 0 diff --git a/tests/08_EXX/53_GO_LR_HF/result.ref b/tests/08_EXX/53_GO_LR_HF/result.ref index 4f161b0cd1..dec386ff49 100644 --- a/tests/08_EXX/53_GO_LR_HF/result.ref +++ b/tests/08_EXX/53_GO_LR_HF/result.ref @@ -1,2 +1,5 @@ -totexcitationenergyref 5.646360 -totaltimeref 3.13 +excitationenergyref1 1.500140 +excitationenergyref2 1.500670 +excitationenergyref3 1.322310 +excitationenergyref4 1.323240 +totaltimeref 1.94 diff --git a/tests/08_EXX/54_GO_ULR_HF/KPT b/tests/08_EXX/54_GO_ULR_HF/KPT new file mode 100644 index 0000000000..c289c0158a --- /dev/null +++ b/tests/08_EXX/54_GO_ULR_HF/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 1 0 0 0 diff --git a/tests/08_EXX/54_GO_ULR_HF/result.ref b/tests/08_EXX/54_GO_ULR_HF/result.ref index de0d1aa940..e8aae3c706 100644 --- a/tests/08_EXX/54_GO_ULR_HF/result.ref +++ b/tests/08_EXX/54_GO_ULR_HF/result.ref @@ -1,2 +1,4 @@ -totexcitationenergyref 0.0 -totaltimeref 3.12 +excitationenergyref1 -0.981255 +excitationenergyref2 -0.977372 +excitationenergyref3 -0.765054 +totaltimeref 1.77 diff --git a/tests/08_EXX/55_KP_LR/result.ref b/tests/08_EXX/55_KP_LR/result.ref index f8872dac9b..68181039ad 100644 --- a/tests/08_EXX/55_KP_LR/result.ref +++ b/tests/08_EXX/55_KP_LR/result.ref @@ -1,2 +1,5 @@ -totexcitationenergyref 0.756248 -totaltimeref 2.24 +excitationenergyref1 0.193805 +excitationenergyref2 0.198422 +excitationenergyref3 0.182005 +excitationenergyref4 0.182016 +totaltimeref 2.25 diff --git a/tests/08_EXX/56_GO_RPA_OUTPUT/result.ref b/tests/08_EXX/56_GO_RPA_OUTPUT/result.ref index 1314357930..7d9ed87963 100644 --- a/tests/08_EXX/56_GO_RPA_OUTPUT/result.ref +++ b/tests/08_EXX/56_GO_RPA_OUTPUT/result.ref @@ -1,16 +1,15 @@ -etotref -422.0200274432613 -etotperatomref -140.6733424811 -Etot_without_rpa -15.238164146732807 -CompareRPA_Cs_data_0_txt_pass 0 -CompareRPA_Cs_data_1_txt_pass 0 -CompareRPA_Cs_data_2_txt_pass 0 -CompareRPA_Cs_data_3_txt_pass 0 -CompareRPA_coulomb_cut_0_txt_pass 0 -CompareRPA_coulomb_cut_1_txt_pass 0 -CompareRPA_coulomb_cut_2_txt_pass 0 -CompareRPA_coulomb_cut_3_txt_pass 0 -CompareRPA_coulomb_mat_0_txt_pass 0 -CompareRPA_coulomb_mat_1_txt_pass 0 -CompareRPA_coulomb_mat_2_txt_pass 0 -CompareRPA_coulomb_mat_3_txt_pass 0 -totaltimeref 1.28 +etotref -15.508944393858402 +Etot_without_rpa -15.238164146732785 +CompareRPA_Cs_data_0_txt_pass 1 +CompareRPA_Cs_data_1_txt_pass 1 +CompareRPA_Cs_data_2_txt_pass 1 +CompareRPA_Cs_data_3_txt_pass 1 +CompareRPA_coulomb_cut_0_txt_pass 1 +CompareRPA_coulomb_cut_1_txt_pass 1 +CompareRPA_coulomb_cut_2_txt_pass 1 +CompareRPA_coulomb_cut_3_txt_pass 1 +CompareRPA_coulomb_mat_0_txt_pass 1 +CompareRPA_coulomb_mat_1_txt_pass 1 +CompareRPA_coulomb_mat_2_txt_pass 1 +CompareRPA_coulomb_mat_3_txt_pass 1 +totaltimeref 1.06 diff --git a/tests/08_EXX/57_KP_RPA_SHRINK/result.ref b/tests/08_EXX/57_KP_RPA_SHRINK/result.ref index 5aa3e241e5..4c97b1624a 100644 --- a/tests/08_EXX/57_KP_RPA_SHRINK/result.ref +++ b/tests/08_EXX/57_KP_RPA_SHRINK/result.ref @@ -1,24 +1,23 @@ -etotref -34.21719023086459 -etotperatomref -17.1085951154 -Etot_without_rpa -1.228213816470492 -CompareRPA_Cs_data_0_txt_pass 0 -CompareRPA_Cs_data_1_txt_pass 0 -CompareRPA_Cs_data_2_txt_pass 0 -CompareRPA_Cs_data_3_txt_pass 0 -CompareRPA_Cs_shrinked_data_0_txt_pass 0 -CompareRPA_Cs_shrinked_data_1_txt_pass 0 -CompareRPA_Cs_shrinked_data_2_txt_pass 0 -CompareRPA_Cs_shrinked_data_3_txt_pass 0 -CompareRPA_coulomb_cut_0_txt_pass 0 -CompareRPA_coulomb_cut_1_txt_pass 0 -CompareRPA_coulomb_cut_2_txt_pass 0 -CompareRPA_coulomb_cut_3_txt_pass 0 -CompareRPA_coulomb_mat_0_txt_pass 0 -CompareRPA_coulomb_mat_1_txt_pass 0 -CompareRPA_coulomb_mat_2_txt_pass 0 -CompareRPA_coulomb_mat_3_txt_pass 0 -CompareRPA_shrink_sinvS_0_txt_pass 0 -CompareRPA_shrink_sinvS_1_txt_pass 0 -CompareRPA_shrink_sinvS_2_txt_pass 0 -CompareRPA_shrink_sinvS_3_txt_pass 0 -totaltimeref 0.75 +etotref -1.257458096999657 +Etot_without_rpa -1.228213816470960 +CompareRPA_Cs_data_0_txt_pass 1 +CompareRPA_Cs_data_1_txt_pass 1 +CompareRPA_Cs_data_2_txt_pass 1 +CompareRPA_Cs_data_3_txt_pass 1 +CompareRPA_Cs_shrinked_data_0_txt_pass 1 +CompareRPA_Cs_shrinked_data_1_txt_pass 1 +CompareRPA_Cs_shrinked_data_2_txt_pass 1 +CompareRPA_Cs_shrinked_data_3_txt_pass 1 +CompareRPA_coulomb_cut_0_txt_pass 1 +CompareRPA_coulomb_cut_1_txt_pass 1 +CompareRPA_coulomb_cut_2_txt_pass 1 +CompareRPA_coulomb_cut_3_txt_pass 1 +CompareRPA_coulomb_mat_0_txt_pass 1 +CompareRPA_coulomb_mat_1_txt_pass 1 +CompareRPA_coulomb_mat_2_txt_pass 1 +CompareRPA_coulomb_mat_3_txt_pass 1 +CompareRPA_shrink_sinvS_0_txt_pass 1 +CompareRPA_shrink_sinvS_1_txt_pass 1 +CompareRPA_shrink_sinvS_2_txt_pass 1 +CompareRPA_shrink_sinvS_3_txt_pass 1 +totaltimeref 0.45 diff --git a/tests/integrate/tools/catch_properties.sh b/tests/integrate/tools/catch_properties.sh index 9acaa6a71c..fbd9670740 100755 --- a/tests/integrate/tools/catch_properties.sh +++ b/tests/integrate/tools/catch_properties.sh @@ -767,13 +767,13 @@ fi # Linear response function #-------------------------------------------- if [ $is_lr == 1 ]; then - lrns=$(get_input_key_value "lr_nstates" "INPUT") - lrns1=`echo "$lrns + 1" |bc` - grep -A$lrns1 "Excitation Energy" $running_path | awk 'NR > 2 && $2 ~ /^[0-9]+\.[0-9]+$/ {print $2}' > lr_eig.txt - lreig_tot=`sum_file lr_eig.txt` - echo "totexcitationenergyref $lreig_tot" >>$1 + shopt -s nullglob + lr_files=(OUT.autotest/trans_analysis_*_tda.dat) + if [ ${#lr_files[@]} -gt 0 ]; then + cat "${lr_files[@]}" | awk '/Excitation Energy/{p=1; next} p && /^[[:space:]]*[0-9]+[[:space:]]/{printf "excitationenergyref%d %.6f\n", ++n, $2} /Occupied orbital/{p=0}' >>$1 + fi + shopt -u nullglob fi - #-------------------------------------------- # Check RDMFT method #-------------------------------------------- diff --git a/tools/02_postprocessing/plot-tools/plot_cond_silce.py b/tools/02_postprocessing/plot-tools/plot_cond_silce.py new file mode 100644 index 0000000000..71b815b963 --- /dev/null +++ b/tools/02_postprocessing/plot-tools/plot_cond_silce.py @@ -0,0 +1,385 @@ +#!/usr/bin/env python3 +"""Plot an ABACUS exciton-density slice. + +The current text format is deliberately compact and self-describing:: + + # ABACUS_EXCITON_SLICE 1 + # state 0 energy_Ry 0.2265 density_kind conditional_elec + # fixed_particle hole position_bohr 2.5 2.5 2.5 + # plane ca slice_pos_bohr 0 + # grid 199 199 u_range -5 6 v_range -5 6 + # bvk 7 7 + # lattice_bohr + # a 0 5.15 5.15 + # b 5.15 0 5.15 + # c 5.15 5.15 0 + # axes_bohr + # u 5.15 5.15 0 + # v 0 5.15 5.15 + # atoms 2 + # Si 0 0 0 + # Si 2.575 2.575 2.575 + # data + + +Files written before format version 1 are still accepted. + +Usage: + python plot_cond_silce.py FILE.dat [-o FILE.png] +""" + +import argparse +from pathlib import Path + +import numpy as np + + +_MAGIC = "ABACUS_EXCITON_SLICE" + + +def _vector(values): + return np.asarray([float(value) for value in values], dtype=float) + + +def _parse_v1(comments): + if comments[0].split() != [_MAGIC, "1"]: + raise ValueError("Unsupported exciton slice format: " + comments[0]) + + meta = {"format_version": 1, "atoms": []} + atom_lines_remaining = 0 + for line in comments[1:]: + fields = line.split() + if not fields: + continue + if atom_lines_remaining: + if len(fields) != 4: + raise ValueError("Invalid atom record: " + line) + meta["atoms"].append({"label": fields[0], "pos": _vector(fields[1:])}) + atom_lines_remaining -= 1 + continue + + key = fields[0] + if key == "state": + if len(fields) != 6 or fields[2] != "energy_Ry" or fields[4] != "density_kind": + raise ValueError("Invalid state record: " + line) + meta["state"] = int(fields[1]) + meta["energy_Ry"] = float(fields[3]) + meta["density_kind"] = fields[5] + elif key == "fixed_particle": + meta["fixed_particle"] = fields[1] + if len(fields) == 6 and fields[2] == "position_bohr": + meta["hole"] = _vector(fields[3:6]) + elif len(fields) != 2: + raise ValueError("Invalid fixed-particle record: " + line) + elif key == "plane": + if len(fields) != 4 or fields[2] != "slice_pos_bohr": + raise ValueError("Invalid plane record: " + line) + meta["plane"] = fields[1] + meta["slice_pos"] = float(fields[3]) + elif key == "grid": + if (len(fields) != 9 or fields[3] != "u_range" + or fields[6] != "v_range"): + raise ValueError("Invalid grid record: " + line) + meta["nu"], meta["nv"] = int(fields[1]), int(fields[2]) + meta["u_min"], meta["u_max"] = float(fields[4]), float(fields[5]) + meta["v_min"], meta["v_max"] = float(fields[7]), float(fields[8]) + elif key == "bvk": + if len(fields) != 3: + raise ValueError("Invalid BvK record: " + line) + meta["nk_u"], meta["nk_v"] = int(fields[1]), int(fields[2]) + elif key in ("a", "b", "c", "u", "v"): + if len(fields) != 4: + raise ValueError("Invalid vector record: " + line) + meta[key if key in ("a", "b", "c") else key + "_vec"] = _vector(fields[1:]) + elif key == "atoms": + if len(fields) != 2: + raise ValueError("Invalid atoms record: " + line) + meta["n_atoms"] = int(fields[1]) + atom_lines_remaining = meta["n_atoms"] + elif key not in ("lattice_bohr", "axes_bohr", "data"): + raise ValueError("Unknown format-v1 header record: " + line) + + if atom_lines_remaining: + raise ValueError("Slice header ended before all atom records were read") + return meta + + +def _parse_legacy(comments): + raw = {} + for line in comments: + fields = line.split(None, 1) + if len(fields) == 2: + raw[fields[0]] = fields[1] + + integer_keys = ( + "state", "grid_nu", "grid_nv", "BvK_u", "BvK_v", + "u_min", "u_max", "v_min", "v_max", "n_atoms", "has_hole", + ) + for key in integer_keys: + if key in raw: + raw[key] = int(raw[key]) + + float_keys = ( + "energy_Ry", "slice_pos", "hole_fix_x", "hole_fix_y", "hole_fix_z", + "u_vec_x", "u_vec_y", "u_vec_z", "v_vec_x", "v_vec_y", "v_vec_z", + "cell_a_x", "cell_a_y", "cell_a_z", "cell_b_x", "cell_b_y", "cell_b_z", + "cell_c_x", "cell_c_y", "cell_c_z", + ) + for key in float_keys: + if key in raw: + raw[key] = float(raw[key]) + + meta = { + "format_version": 0, + "state": raw["state"], + "energy_Ry": raw["energy_Ry"], + "density_kind": raw.get("density_kind", "conditional_elec"), + "fixed_particle": raw.get("fixed_particle", "hole"), + "plane": raw.get("plane", "ab"), + "slice_pos": raw.get("slice_pos", 0.0), + "nu": raw["grid_nu"], + "nv": raw["grid_nv"], + "nk_u": raw.get("BvK_u", 1), + "nk_v": raw.get("BvK_v", 1), + "u_min": raw["u_min"], + "u_max": raw["u_max"], + "v_min": raw["v_min"], + "v_max": raw["v_max"], + "u_vec": _vector([raw["u_vec_x"], raw["u_vec_y"], raw["u_vec_z"]]), + "v_vec": _vector([raw["v_vec_x"], raw["v_vec_y"], raw["v_vec_z"]]), + "a": _vector([raw["cell_a_x"], raw["cell_a_y"], raw["cell_a_z"]]), + "b": _vector([raw["cell_b_x"], raw["cell_b_y"], raw["cell_b_z"]]), + "c": _vector([raw["cell_c_x"], raw["cell_c_y"], raw["cell_c_z"]]), + "atoms": [], + } + has_fixed = raw.get( + "has_hole", 1 if meta["density_kind"].startswith("conditional") else 0) + if has_fixed: + meta["hole"] = _vector([ + raw["hole_fix_x"], raw["hole_fix_y"], raw["hole_fix_z"]]) + for index in range(raw["n_atoms"]): + meta["atoms"].append({ + "label": raw["atom_label_{}".format(index)], + "pos": _vector([ + raw["atom_x_{}".format(index)], + raw["atom_y_{}".format(index)], + raw["atom_z_{}".format(index)], + ]), + }) + meta["n_atoms"] = len(meta["atoms"]) + return meta + + +def _validate(meta, data, filename): + required = ( + "state", "energy_Ry", "density_kind", "fixed_particle", "plane", + "slice_pos", "nu", "nv", "nk_u", "nk_v", "u_min", "u_max", + "v_min", "v_max", "u_vec", "v_vec", "a", "b", "c", "atoms", + ) + missing = [key for key in required if key not in meta] + if missing: + raise ValueError("{} is missing metadata: {}".format( + filename, ", ".join(missing))) + if meta["plane"] not in ("ab", "bc", "ca"): + raise ValueError("Unsupported slice plane: " + meta["plane"]) + if data.shape != (meta["nu"], meta["nv"]): + raise ValueError( + "{} declares a {}x{} grid but contains {}".format( + filename, meta["nu"], meta["nv"], data.shape)) + if not np.all(np.isfinite(data)): + raise ValueError("{} contains non-finite density values".format(filename)) + if np.any(data < 0.0): + raise ValueError("{} contains negative density values".format(filename)) + if np.linalg.norm(meta["u_vec"]) == 0 or np.linalg.norm(meta["v_vec"]) == 0: + raise ValueError("{} contains a zero slice vector".format(filename)) + + +def read_slice(filename): + """Read metadata and the density matrix from an old or format-v1 file.""" + comments = [] + with open(filename, encoding="utf-8") as stream: + for raw_line in stream: + stripped = raw_line.strip() + if stripped.startswith("#"): + comments.append(stripped[1:].strip()) + if not comments: + raise ValueError("{} has no exciton slice header".format(filename)) + + meta = (_parse_v1(comments) if comments[0].startswith(_MAGIC) + else _parse_legacy(comments)) + data = np.loadtxt(filename, comments="#", dtype=float) + data = np.atleast_2d(data) + _validate(meta, data, filename) + return meta, data + + +def parse(filename): + """Compatibility wrapper returning only metadata.""" + return read_slice(filename)[0] + + +def load_data(filename): + """Compatibility wrapper returning only the density matrix.""" + return read_slice(filename)[1] + + +def _density_labels(kind): + labels = { + "conditional_elec": (r"$|\psi_{\rm cond}|^2$", "Conditional Electron Density"), + "conditional_hole": (r"$|\psi_{\rm cond}|^2$", "Conditional Hole Density"), + "average_hole": (r"$\rho_{\rm h}^{\rm avg}$", "Average Hole Density"), + "average_elec": (r"$\rho_{\rm e}^{\rm avg}$", "Average Electron Density"), + } + return labels.get(kind, (kind, kind.replace("_", " ").title())) + + +def plot(data, meta, output): + """Render the density in Cartesian coordinates for a possibly skew plane.""" + import matplotlib.pyplot as plt + from matplotlib.colors import LogNorm + from matplotlib.lines import Line2D + + u_vec, v_vec = meta["u_vec"], meta["v_vec"] + u_length = np.linalg.norm(u_vec) + v_length = np.linalg.norm(v_vec) + u_hat = u_vec / u_length + v_perp = v_vec - np.dot(v_vec, u_hat) * u_hat + v_perp_length = np.linalg.norm(v_perp) + if v_perp_length == 0: + raise ValueError("Slice vectors are collinear") + v_perp_hat = v_perp / v_perp_length + v_along_u = np.dot(v_vec, u_hat) + angle = np.degrees(np.arccos(np.clip( + np.dot(u_vec, v_vec) / (u_length * v_length), -1.0, 1.0))) + + u_edges = np.linspace(meta["u_min"], meta["u_max"], meta["nu"]) + v_edges = np.linspace(meta["v_min"], meta["v_max"], meta["nv"]) + du = u_edges[1] - u_edges[0] + dv = v_edges[1] - v_edges[0] + u_edges = np.r_[u_edges - 0.5 * du, u_edges[-1] + 0.5 * du] + v_edges = np.r_[v_edges - 0.5 * dv, v_edges[-1] + 0.5 * dv] + x_edges = u_edges[:, None] * u_length + v_edges[None, :] * v_along_u + y_edges = np.broadcast_to( + v_edges[None, :] * v_perp_length, x_edges.shape) + + maximum = data.max() + if maximum <= 0: + raise ValueError("Density data contains no positive values") + positive = data[data > 0] + minimum = max(positive.min(), maximum * 1.0e-6) + + fig, axis = plt.subplots(figsize=(12, 10)) + image = axis.pcolormesh( + x_edges, y_edges, data.T, cmap="hot", + norm=LogNorm(vmin=minimum, vmax=maximum), + shading="flat", rasterized=True) + color_label, title_kind = _density_labels(meta["density_kind"]) + fig.colorbar(image, ax=axis, label=color_label, shrink=0.78) + + def project(position): + return np.dot(position, u_hat), np.dot(position, v_perp_hat) + + for index in range(int(np.floor(meta["u_min"])), + int(np.ceil(meta["u_max"])) + 1): + start = index * u_vec + meta["v_min"] * v_vec + end = index * u_vec + meta["v_max"] * v_vec + axis.plot(*zip(project(start), project(end)), + color="black", ls="--", lw=0.6, alpha=0.55) + for index in range(int(np.floor(meta["v_min"])), + int(np.ceil(meta["v_max"])) + 1): + start = meta["u_min"] * u_vec + index * v_vec + end = meta["u_max"] * u_vec + index * v_vec + axis.plot(*zip(project(start), project(end)), + color="black", ls="--", lw=0.6, alpha=0.55) + + bvk_u0 = -(meta["nk_u"] // 2) + bvk_v0 = -(meta["nk_v"] // 2) + for index in (bvk_u0, bvk_u0 + meta["nk_u"]): + axis.plot(*zip( + project(index * u_vec + meta["v_min"] * v_vec), + project(index * u_vec + meta["v_max"] * v_vec)), + color="black", lw=2.8, alpha=0.94) + for index in (bvk_v0, bvk_v0 + meta["nk_v"]): + axis.plot(*zip( + project(meta["u_min"] * u_vec + index * v_vec), + project(meta["u_max"] * u_vec + index * v_vec)), + color="black", lw=2.8, alpha=0.94) + + perpendicular = ( + meta["c"] if meta["plane"] == "ab" + else meta["a"] if meta["plane"] == "bc" else meta["b"]) + perpendicular_length = np.linalg.norm(perpendicular) + perpendicular_hat = perpendicular / perpendicular_length + labels = sorted({atom["label"] for atom in meta["atoms"]}) + colors = plt.cm.tab10(np.linspace(0, 1, max(1, len(labels)))) + label_colors = dict(zip(labels, colors)) + for atom in meta["atoms"]: + home_x, home_y = project(atom["pos"]) + depth = np.dot(atom["pos"], perpendicular_hat) - meta["slice_pos"] + depth = ((depth + perpendicular_length / 2) % perpendicular_length + - perpendicular_length / 2) + alpha = max(0.2, 1.0 - abs(depth) / (perpendicular_length * 0.5)) + for cell_u in range(int(np.floor(meta["u_min"])), + int(np.ceil(meta["u_max"]))): + for cell_v in range(int(np.floor(meta["v_min"])), + int(np.ceil(meta["v_max"]))): + axis.plot( + home_x + cell_u * u_length + cell_v * v_along_u, + home_y + cell_v * v_perp_length, + "o", color=label_colors[atom["label"]], ms=6, + alpha=alpha, mec="white", mew=0.5) + + handles = [] + if "hole" in meta: + fixed_x, fixed_y = project(meta["hole"]) + axis.plot( + fixed_x, fixed_y, "o", color="lime", ms=10, + mec="black", mew=1.5, zorder=10) + handles.append(Line2D( + [0], [0], marker="o", color="white", markerfacecolor="lime", + markersize=8, markeredgecolor="black", markeredgewidth=1.5, + label="Fixed {} (home cell)".format(meta["fixed_particle"]))) + for label in labels: + handles.append(Line2D( + [0], [0], marker="o", color="white", + markerfacecolor=label_colors[label], markersize=7, + markeredgecolor="white", markeredgewidth=0.5, label=label)) + if handles: + axis.legend(loc="upper right", fontsize=9, handles=handles) + + axis.set_xlabel("Distance along u (|u| = {:.2f} Bohr)".format(u_length)) + axis.set_ylabel( + "Distance perpendicular to u (|v| = {:.2f} Bohr, angle = {:.0f} deg)" + .format(v_length, angle)) + axis.set_aspect("equal") + fixed_text = ( + "{} at ({:.2f}, {:.2f}, {:.2f}) Bohr | ".format( + meta["fixed_particle"].title(), *meta["hole"]) + if "hole" in meta else "") + axis.set_title( + "{} | State {} | {}-plane\n{}E = {:.6f} Ry = {:.3f} eV\n" + "BvK {}x{} | grid {}x{}".format( + title_kind, meta["state"], meta["plane"], fixed_text, + meta["energy_Ry"], meta["energy_Ry"] * 13.605693, + meta["nk_u"], meta["nk_v"], meta["nu"], meta["nv"])) + fig.tight_layout() + fig.savefig(output, dpi=150, bbox_inches="tight") + plt.close(fig) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("data_file", type=Path) + parser.add_argument("-o", "--output", type=Path) + args = parser.parse_args() + + meta, data = read_slice(args.data_file) + output = args.output or args.data_file.with_suffix(".png") + plot(data, meta, output) + print("{} | format v{} | grid {} | BvK {}x{} | saved {}".format( + meta["density_kind"], meta["format_version"], data.shape, + meta["nk_u"], meta["nk_v"], output)) + + +if __name__ == "__main__": + main() From bba2f0c31352642b981ff61ee0a9fc98c0bcbd55 Mon Sep 17 00:00:00 2001 From: Fisherd99 Date: Fri, 31 Jul 2026 18:49:25 +0800 Subject: [PATCH 2/2] tmp --- source/source_esolver/esolver.cpp | 23 ++- .../module_lr/bse/esolver_bse_lcao.cpp | 30 +-- .../module_lr/bse/esolver_bse_lcao.h | 4 +- .../source_lcao/module_lr/bse/hamilt_bse.cpp | 2 +- source/source_lcao/module_lr/bse/hamilt_bse.h | 2 +- .../module_lr/esolver_lrtd_lcao.cpp | 191 ++++++++++++------ .../source_lcao/module_lr/esolver_lrtd_lcao.h | 13 +- source/source_lcao/module_lr/hamilt_casida.h | 2 +- .../module_lr/utils/exciton_plotter.h | 2 +- .../module_lr/utils/spectrum_mo.hpp | 2 +- 10 files changed, 179 insertions(+), 92 deletions(-) diff --git a/source/source_esolver/esolver.cpp b/source/source_esolver/esolver.cpp index 036b8abf09..f72adba244 100644 --- a/source/source_esolver/esolver.cpp +++ b/source/source_esolver/esolver.cpp @@ -12,6 +12,7 @@ #include "esolver_ks_lcao_tddft.h" #include "esolver_ks_lcaopw.h" #include "source_lcao/module_lr/esolver_lrtd_lcao.h" +#include "source_lcao/module_lr/bse/esolver_bse_lcao.h" #include "source_base/module_external/blacs_connector.h" #endif #include "esolver_dp.h" @@ -273,13 +274,27 @@ ESolver* init_esolver(const Input_para& inp) } else if (esolver_type == "lr_lcao") { - if (PARAM.globalv.gamma_only_local) + if (PARAM.inp.xc_kernel != "bse") { - return new LR::ESolver_LR(inp); + if (PARAM.globalv.gamma_only_local) + { + return new LR::ESolver_LR(inp); + } + else + { + return new LR::ESolver_LR, double>(inp); + } } else - { - return new LR::ESolver_LR, double>(inp); + { + if (PARAM.globalv.gamma_only_local) + { + return new BSE::ESolver_BSE(inp); + } + else + { + return new BSE::ESolver_BSE, double>(inp); + } } } else if (esolver_type == "ksdft_lr_lcao") diff --git a/source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp b/source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp index c83bbb2e67..37b1e049b8 100644 --- a/source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp +++ b/source/source_lcao/module_lr/bse/esolver_bse_lcao.cpp @@ -2,15 +2,15 @@ #include #include "source_cell/module_neighbor/sltk_atom_arrange.h" #include "source_io/module_output/print_info.h" -#include "source_lcao/module_gint/gint.h" +#include "source_hamilt/module_gint/gint.h" namespace BSE { template -ESolver_BSE::ESolver_BSE(const Input_para& inp, UnitCell& ucell) : - LR::ESolver_LR(inp, ucell, LR::DeferredESolverLRInit{}) +void ESolver_BSE::before_all_runners(UnitCell& ucell, const Input_para& inp) { - ModuleBase::TITLE("ESolver_BSE", "ESolver_BSE(from scratch)"); - ModuleBase::timer::start("ESolver_BSE", "constructor"); + ModuleBase::TITLE("ESolver_BSE", "before_all_runners"); + ModuleBase::timer::start("ESolver_BSE", "before_all_runners"); + this->ucell_ = &ucell; // xc kernel this->xc_kernel = LR_Util::tolower(inp.xc_kernel); @@ -18,7 +18,7 @@ ESolver_BSE::ESolver_BSE(const Input_para& inp, UnitCell& ucell) : ModuleESolver::ESolver_FP::before_all_runners(ucell, inp); this->pelec = new elecstate::ElecStateLCAO(); - this->kRlist = LR_IO::RI_kRlist(this->ucell, &this->kv); + this->kRlist = LR_IO::RI_kRlist(*this->ucell_, &this->kv); ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "Set K-POINTS and R-list for RI"); ModuleIO::print_parameters(ucell, this->kv, inp); @@ -89,7 +89,7 @@ ESolver_BSE::ESolver_BSE(const Input_para& inp, UnitCell& ucell) : atom_arrange::search(PARAM.globalv.search_pbc, GlobalV::ofs_running, this->gd, - this->ucell, + *this->ucell_, search_radius, PARAM.inp.test_atom_input); @@ -114,7 +114,7 @@ ESolver_BSE::ESolver_BSE(const Input_para& inp, UnitCell& ucell) : this->pot.resize(this->nspin, nullptr); if (this->input.lr_solver != "spectrum" && this->input.lr_solver != "plot") { - this->mo_lri = LR_Util::make_unique>(this->ucell, + this->mo_lri = LR_Util::make_unique>(*this->ucell_, this->nk, this->kRlist, this->nocc[0], @@ -128,7 +128,7 @@ ESolver_BSE::ESolver_BSE(const Input_para& inp, UnitCell& ucell) : this->init_pot(chg_gs); } } - ModuleBase::timer::end("ESolver_BSE", "constructor"); + ModuleBase::timer::end("ESolver_BSE", "before_all_runners"); } template @@ -159,7 +159,7 @@ void ESolver_BSE::runner(UnitCell& ucell, const int istep) std::cout << "Calculating Casida/BSE matrix directly." << std::endl; assert(this->xc_kernel == "bse"); this->lri_init(); - HamiltBSE bse_matrix(this->nspin, this->nbasis, this->nocc, this->nvirt, this->ucell, + HamiltBSE bse_matrix(this->nspin, this->nbasis, this->nocc, this->nvirt, *this->ucell_, this->orb_cutoff_, this->gd, *this->psi_ks, *this->psi_ks_global, this->eig_gw, *this->mo_lri, this->pot[0], this->kv, this->paraX_, this->paraC_, this->paraMat_, @@ -358,7 +358,7 @@ void ESolver_BSE::after_all_runners(UnitCell& ucell) std::cout << "plot BSE exciton wavefunction for state: " << this->input.plot_istate << ", spin type: " << this->input.bse_spin_types[is] << std::endl; LR_Util::ExcitonPlotter eplot(this->nspin, this->nbasis, this->nocc, this->nvirt, *this->psi_ks, - this->ucell, this->kv, this->gd, this->orb_cutoff_, this->Pgrid, *this->pw_rho, + *this->ucell_, this->kv, this->gd, this->orb_cutoff_, this->Pgrid, *this->pw_rho, this->paraX_, this->paraC_, this->paraMat_, output_dir, &this->tda_ene[is * this->nstates], this->X[is].template data(), @@ -441,7 +441,7 @@ void ESolver_BSE::after_all_runners(UnitCell& ucell) if (LR_Util::tolower(this->input.abs_gauge) == "velocity" ) { const int nspin_tmp = PARAM.inp.nspin == 2 ? 2 : 1; - this->velocity_mo = LR_Util::cal_velocity_mo(this->ucell, this->gd, this->two_center_bundle_, + this->velocity_mo = LR_Util::cal_velocity_mo(*this->ucell_, this->gd, this->two_center_bundle_, this->paraMat_, this->paraC_, this->kv, *this->psi_ks, this->nk, nspin_tmp, this->nbasis, this->nocc, this->nvirt); } @@ -450,7 +450,7 @@ void ESolver_BSE::after_all_runners(UnitCell& ucell) for (int is = 0; is < this->X.size(); ++is) { LR::LR_Spectrum spectrum(this->nspin, this->nbasis, this->nocc, this->nvirt, *this->pw_rho, *this->psi_ks, - this->ucell, this->kv, this->gd, this->orb_cutoff_, this->two_center_bundle_, + *this->ucell_, this->kv, this->gd, this->orb_cutoff_, this->two_center_bundle_, this->paraX_, this->paraC_, this->paraMat_, &this->tda_ene[is * this->nstates], this->eig_ks.c, this->X[is].template data(), this->nstates, false/*openshell*/, @@ -481,7 +481,7 @@ void ESolver_BSE::after_all_runners(UnitCell& ucell) for (int is = 0;is < this->full_X.size();++is) { LR::LR_Spectrum spectrum(this->nspin, this->nbasis, this->nocc, this->nvirt, *this->pw_rho, *this->psi_ks, - this->ucell, this->kv, this->gd, this->orb_cutoff_, this->two_center_bundle_, + *this->ucell_, this->kv, this->gd, this->orb_cutoff_, this->two_center_bundle_, this->paraX_, this->paraC_, this->paraMat_, &this->full_ene[is * this->nstates], this->eig_ks.c, this->full_X[is].template data(), this->nstates, false/*openshell*/, @@ -672,7 +672,7 @@ void ESolver_BSE::init_pot(const Charge& chg_gs) { using ST = LR::PotHxcLR::SpinType; case 1: case 2: - this->pot[0] = std::make_shared(this->xc_kernel, *this->pw_rho, this->ucell, chg_gs, this->Pgrid, + this->pot[0] = std::make_shared(this->xc_kernel, *this->pw_rho, *this->ucell_, chg_gs, this->Pgrid, ST::S1, this->input.lr_init_xc_kernel); break; // case 2: diff --git a/source/source_lcao/module_lr/bse/esolver_bse_lcao.h b/source/source_lcao/module_lr/bse/esolver_bse_lcao.h index 52767f7555..4fb8ef3758 100644 --- a/source/source_lcao/module_lr/bse/esolver_bse_lcao.h +++ b/source/source_lcao/module_lr/bse/esolver_bse_lcao.h @@ -15,8 +15,7 @@ namespace BSE template class ESolver_BSE : public LR::ESolver_LR { public: - /// @brief a from-scratch constructor - ESolver_BSE(const Input_para& inp, UnitCell& ucell); + explicit ESolver_BSE(const Input_para& inp) : LR::ESolver_LR(inp) {} ~ESolver_BSE() override { //delete this->psi_ks; // already deleted in ESolver_LR @@ -35,6 +34,7 @@ namespace BSE std::unique_ptr> mo_lri; + virtual void before_all_runners(UnitCell& ucell, const Input_para& inp) override; virtual void runner(UnitCell& ucell, int istep) override; virtual void after_all_runners(UnitCell& ucell) override; diff --git a/source/source_lcao/module_lr/bse/hamilt_bse.cpp b/source/source_lcao/module_lr/bse/hamilt_bse.cpp index f0a62df7c5..d4dfc069af 100644 --- a/source/source_lcao/module_lr/bse/hamilt_bse.cpp +++ b/source/source_lcao/module_lr/bse/hamilt_bse.cpp @@ -1,7 +1,7 @@ #include "hamilt_bse.h" #include "hamilt_bse_solver.h" -#include "source_lcao/module_gint/gint_interface.h" +#include "source_hamilt/module_gint/gint_interface.h" #include "source_lcao/module_lr/utils/lr_util_hcontainer.h" #include diff --git a/source/source_lcao/module_lr/bse/hamilt_bse.h b/source/source_lcao/module_lr/bse/hamilt_bse.h index 4013a1cf49..ea78ac2a00 100644 --- a/source/source_lcao/module_lr/bse/hamilt_bse.h +++ b/source/source_lcao/module_lr/bse/hamilt_bse.h @@ -13,7 +13,7 @@ #include "source_cell/unitcell.h" #include "source_estate/module_dm/density_matrix.h" #include "source_hamilt/hamilt.h" -#include "source_lcao/module_hcontainer/hcontainer_funcs.h" +#include "source_hamilt/module_hcontainer/hcontainer_funcs.h" #include "source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h" #include "source_lcao/module_lr/potentials/pot_hxc_lrtd.h" #include "source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h" diff --git a/source/source_lcao/module_lr/esolver_lrtd_lcao.cpp b/source/source_lcao/module_lr/esolver_lrtd_lcao.cpp index ab969364ba..315df79979 100644 --- a/source/source_lcao/module_lr/esolver_lrtd_lcao.cpp +++ b/source/source_lcao/module_lr/esolver_lrtd_lcao.cpp @@ -49,7 +49,8 @@ void LR::ESolver_LR::move_exx_lri(std::shared_ptr +int LR::ESolver_LR::cal_nupdown_form_occ(const ModuleBase::matrix& wg) { // only for nspin=2 const int& nk = wg.nr / 2; auto occ_sum_k = [&](const int& is, const int& ib)->double { double o = 0.0; for (int ik = 0;ik < nk;++ik) { o += wg(is * nk + ik, ib); } return o;}; @@ -64,7 +65,8 @@ inline int cal_nupdown_form_occ(const ModuleBase::matrix& wg) return nupdown; } -inline void setup_2center_table(TwoCenterBundle& two_center_bundle, LCAO_Orbitals& orb, UnitCell& ucell) +template +void LR::ESolver_LR::setup_2center_table(TwoCenterBundle& two_center_bundle, LCAO_Orbitals& orb, UnitCell& ucell) { // set up 2-center table #ifdef __FFT_TWO_CENTER @@ -86,17 +88,21 @@ inline void setup_2center_table(TwoCenterBundle& two_center_bundle, LCAO_Orbital template void LR::ESolver_LR::parameter_check()const { - const std::set lr_solvers = { "dav", "lapack" , "spectrum", "dav_subspace", "cg" }; - const std::set xc_kernels = { "rpa", "lda", "pwlda", "pbe", "hf" , "hse" }; + const std::set lr_solvers = { "dav", "lapack" , "spectrum", "dav_subspace", "cg", "elpa", "plot" }; + const std::set xc_kernels = { "rpa", "lda", "pwlda", "pbe", "hf", "hse", "bse" }; + const std::set abs_gauge = { "velocity", "length" }; if (lr_solvers.find(this->input.lr_solver) == lr_solvers.end()) { throw std::invalid_argument("ESolver_LR: unknown type of lr_solver"); -} + } if (xc_kernels.find(this->xc_kernel) == xc_kernels.end()) { throw std::invalid_argument("ESolver_LR: unknown type of xc_kernel"); -} + } if (this->nspin != 1 && this->nspin != 2) { throw std::invalid_argument("LR-TDDFT only supports nspin = 1 or 2 now"); -} + } + if (abs_gauge.find(this->input.abs_gauge) == abs_gauge.end()) { + throw std::invalid_argument("ESolver_LR: unknown type of abs_gauge"); + } } template @@ -105,17 +111,48 @@ void LR::ESolver_LR::set_dimension() this->nspin = PARAM.inp.nspin; this->nstates = input.lr_nstates; this->nbasis = PARAM.globalv.nlocal; - // calculate the number of occupied and unoccupied states - // which determines the basis size of the excited states + int ks_nbands = PARAM.inp.nbands; this->nocc_max = LR_Util::cal_nocc(LR_Util::cal_nelec(*this->ucell_)); + if (input.ri_hartree_benchmark == "aims" || input.ri_hartree_benchmark == "aims-librpa" + && !input.aims_nbasis.empty()) + { + // calculate total number of basis funcs, see https://en.cppreference.com/w/cpp/algorithm/inner_product + this->nbasis = std::inner_product(input.aims_nbasis.begin(), /* iterator1.begin */ + input.aims_nbasis.end(), /* iterator1.end */ + this->ucell_->atoms, /* iterator2.begin */ + 0, /* init value */ + std::plus(), /* iter op1 */ + [](const int& a, const Atom& b) { return a * b.na; }); /* iter op2 */ + std::cout << "nbasis from aims: " << this->nbasis << std::endl; + for (int it = 0; it < this->ucell_->ntype; ++it) + { + this->ucell_->atoms[it].nw = input.aims_nbasis[it]; + } + const_cast(this->ucell_)->set_iat2iwt(1); // update iat2iwt for aims_nbasis 25-05-23 + + int nbands_file = 0; + int nk_file = 0; + int nspin_file = 0; + int nocc_file = 0; + LR_IO::parse_band_out_file(nbands_file, nk_file, nspin_file, nocc_file); + std::cout << "nocc from band_out: " << nocc_file << std::endl; + ks_nbands = nbands_file; + this->nocc_max = nocc_file; + } + // calculate the number of occupied and unoccupied states + // which determines the basis size of the excited states this->nocc_in = std::max(1, std::min(input.nocc, this->nocc_max)); - this->nvirt_in = PARAM.inp.nbands - this->nocc_max; //nbands-nocc - if (input.nvirt > this->nvirt_in) { GlobalV::ofs_warning << "ESolver_LR: input nvirt is too large to cover by nbands, set nvirt = nbands - nocc = " << this->nvirt_in << std::endl; } + this->nvirt_in = ks_nbands - this->nocc_max; //nbands-nocc + if (input.nvirt > this->nvirt_in) { GlobalV::ofs_running << "ESolver_LR: input nvirt is too large to cover by nbands, set nvirt = nbands - nocc = " << this->nvirt_in << std::endl; } else if (input.nvirt > 0) { this->nvirt_in = input.nvirt; } this->nbands = this->nocc_in + this->nvirt_in; - this->nk = this->kv.get_nks() / this->nspin; + this->nk = PARAM.inp.nspin == 2 ? this->kv.get_nks() / 2 : this->kv.get_nks(); this->nocc.resize(nspin, nocc_in); this->nvirt.resize(nspin, nvirt_in); + if (this->nstates <= 0) { + this->nstates = nk * nocc_in * nvirt_in; + GlobalV::ofs_running << "ESolver_LR: lr_nstates <= 0, set nstates = nk * nocc * nvirt = " << this->nstates << std::endl; + } for (int is = 0;is < nspin;++is) { this->npairs.push_back(nocc[is] * nvirt[is]); } GlobalV::ofs_running << "Setting LR-TDDFT parameters: " << std::endl; GlobalV::ofs_running << "number of occupied bands: " << nocc_in << std::endl; @@ -123,17 +160,6 @@ void LR::ESolver_LR::set_dimension() GlobalV::ofs_running << "number of Atom orbitals (LCAO-basis size): " << this->nbasis << std::endl; GlobalV::ofs_running << "number of KS bands: " << this->eig_ks.nc << std::endl; GlobalV::ofs_running << "number of excited states to be solved: " << this->nstates << std::endl; - if (input.ri_hartree_benchmark == "aims" && !input.aims_nbasis.empty()) - { - // calculate total number of basis funcs, see https://en.cppreference.com/w/cpp/algorithm/inner_product - this->nbasis = std::inner_product(input.aims_nbasis.begin(), /* iterator1.begin */ - input.aims_nbasis.end(), /* iterator1.end */ - this->ucell_->atoms, /* iterator2.begin */ - 0, /* init value */ - std::plus(), /* iter op1 */ - [](const int& a, const Atom& b) { return a * b.na; }); /* iter op2 */ - std::cout << "nbasis from aims: " << this->nbasis << std::endl; - } } template @@ -360,6 +386,7 @@ void LR::ESolver_LR::initialize_from_unitcell_(UnitCell& ucell, const Inp #ifdef __MPI this->paraMat_.set_desc_wfc_Eij(this->nbasis, this->nbands, paraMat_.get_row_size()); int err = this->paraMat_.set_nloc_wfc_Eij(this->nbands, GlobalV::ofs_running, GlobalV::ofs_warning); + this->paraMat_.set_atomic_trace(ucell.get_iat2iwt(), ucell.nat, this->nbasis); if (input.ri_hartree_benchmark != "aims") { this->paraMat_.set_atomic_trace(ucell.get_iat2iwt(), ucell.nat, this->nbasis); } #else this->paraMat_.nrow_bands = this->nbasis; @@ -387,7 +414,7 @@ void LR::ESolver_LR::initialize_from_unitcell_(UnitCell& ucell, const Inp #endif ); - //allocate 2-particle state and setup 2d division + // clear ks info, new elecstate for excition this->pelec = new elecstate::ElecState(); // read the ground state charge density and calculate xc kernel @@ -399,7 +426,7 @@ void LR::ESolver_LR::initialize_from_unitcell_(UnitCell& ucell, const Inp pw_big->nbz, pw_big->bz); Charge chg_gs; - if (input.ri_hartree_benchmark != "aims") { this->read_ks_chg(chg_gs); } + if (input.ri_hartree_benchmark == "none") { this->read_ks_chg(chg_gs); } this->init_pot(chg_gs); // search adjacent atoms and init Gint @@ -458,24 +485,31 @@ void LR::ESolver_LR::runner(UnitCell& ucell, const int istep) this->setup_eigenvectors_X(); this->pelec->ekb.create(nspin, this->nstates); - auto efile = [&](const std::string& label)->std::string {return PARAM.globalv.global_out_dir + "Excitation_Energy_" + label + ".dat";}; - auto vfile = [&](const std::string& label)->std::string {return PARAM.globalv.global_out_dir + "Excitation_Amplitude_" + label + "_" + std::to_string(GlobalV::MY_RANK) + ".dat";}; - if (this->input.lr_solver != "spectrum") + auto efile_out = [&](const std::string& label)->std::string {return PARAM.globalv.global_out_dir + "Excitation_Energy_" + label + ".dat";}; + auto vfile_out = [&](const std::string& label)->std::string {return PARAM.globalv.global_out_dir + "Excitation_Amplitude_" + label + "_" + std::to_string(GlobalV::MY_RANK) + ".dat";}; + if (this->input.lr_solver == "elpa") + { + ModuleBase::WARNING_QUIT("ESolver_LR", "ESolver_LR doesn't support elpa now."); + + } + else if (this->input.lr_solver != "spectrum") { auto write_states = [&](const std::string& label, const Real* e, const T* v, const int& dim, const int& nst, const int& prec = 8)->void { - if (GlobalV::MY_RANK == 0) { assert(nst == LR_Util::write_value(efile(label), prec, e, nst)); } - assert(nst * dim == LR_Util::write_value(vfile(label), prec, v, nst, dim)); + if (GlobalV::MY_RANK == 0) { assert(nst == LR_Util::write_value(efile_out(label), prec, e, nst)); } + assert(nst * dim == LR_Util::write_value(vfile_out(label), prec, v, nst, dim)); }; - std::vector precondition(this->input.lr_solver == "lapack" ? 0 : nloc_per_band, 1.0); + std::vector precondition(this->input.lr_solver == "lapack" ? 0 : nloc_per_state, 1.0); // allocate and initialize A matrix and density matrix if (openshell) { for (int is : {0, 1}) { - const int offset_is = is * this->paraX_[0].get_local_size(); - OperatorLRDiag pre_op(this->eig_ks.c + is * nk * (nocc[0] + nvirt[0]), this->paraX_[is], this->nk, this->nocc[is], this->nvirt[is]); - if (input.lr_solver != "lapack") { pre_op.act(1, offset_is, 1, precondition.data() + offset_is, precondition.data() + offset_is); } + if (input.lr_solver != "lapack") { + const int offset_is = is * this->paraX_[0].get_local_size(); + OperatorLRDiag pre_op(this->eig_ks.c + is * nk * (nocc[0] + nvirt[0]), this->paraX_[is], this->nk, this->nocc[is], this->nvirt[is]); + pre_op.act(1, offset_is, 1, precondition.data() + offset_is, precondition.data() + offset_is); + } } std::cout << "Solving spin-conserving excitation for open-shell system." << std::endl; HamiltULR hulr(xc_kernel, @@ -497,13 +531,18 @@ void LR::ESolver_LR::runner(UnitCell& ucell, const int istep) this->paraX_, this->paraC_, this->paraMat_); - LR::HSolver::solve(hulr, this->X[0].template data(), nloc_per_band, nstates, this->pelec->ekb.c, this->input.lr_solver, this->input.lr_thr, precondition); - if (input.out_wfc_lr) { write_states("openshell", this->pelec->ekb.c, this->X[0].template data(), nloc_per_band, nstates); } + LR::HSolver::solve(hulr, this->X[0].template data(), nloc_per_state, nstates, + this->nk, this->nocc, this->nvirt, this->paraX_, + this->pelec->ekb.c, this->input.lr_solver, + this->input.lr_thr, precondition); + if (input.out_wfc_lr) { write_states("openshell", this->pelec->ekb.c, this->X[0].template data(), nloc_per_state, nstates); } } else { - OperatorLRDiag pre_op(this->eig_ks.c, this->paraX_[0], this->nk, this->nocc[0], this->nvirt[0]); - if (input.lr_solver != "lapack") { pre_op.act(1, nloc_per_band, 1, precondition.data(), precondition.data()); } + if (input.lr_solver != "lapack") { + OperatorLRDiag pre_op(this->eig_ks.c, this->paraX_[0], this->nk, this->nocc[0], this->nvirt[0]); + pre_op.act(1, nloc_per_state, 1, precondition.data(), precondition.data()); + } auto spin_types = std::vector({ "singlet", "triplet" }); for (int is = 0;is < nspin;++is) { @@ -530,30 +569,43 @@ void LR::ESolver_LR::runner(UnitCell& ucell, const int istep) spin_types[is], input.ri_hartree_benchmark, (input.ri_hartree_benchmark == "aims" ? input.aims_nbasis : std::vector({}))); - // solve the Casida equation - LR::HSolver::solve(hlr, this->X[is].template data(), nloc_per_band, nstates, - this->pelec->ekb.c + is * nstates, this->input.lr_solver, this->input.lr_thr, precondition/*, - !std::set({ "hf", "hse" }).count(this->xc_kernel)*/); //whether the kernel is Hermitian - if (input.out_wfc_lr) { write_states(spin_types[is], this->pelec->ekb.c + is * nstates, this->X[is].template data(), nloc_per_band, nstates); } + LR::HSolver::solve(hlr, this->X[is].template data(), nloc_per_state, nstates, + this->nk, this->nocc, this->nvirt, this->paraX_, + this->pelec->ekb.c + is * nstates, + this->input.lr_solver, + this->input.lr_thr, + precondition); + if (input.out_wfc_lr) { write_states(spin_types[is], this->pelec->ekb.c + is * nstates, this->X[is].template data(), nloc_per_state, nstates); } } } } - else // read the eigenvalues + else // lr_solver == "spectrum", read the eigenvalues { + auto efile_in = [&](const std::string& label)->std::string {return PARAM.globalv.global_readin_dir + "Excitation_Energy_" + label + ".dat";}; + auto vfile_in = [&](const std::string& label)->std::string {return PARAM.globalv.global_readin_dir + "Excitation_Amplitude_" + label + "_" + std::to_string(GlobalV::MY_RANK) + ".dat";}; + auto read_states = [&](const std::string& label, Real* e, T* v, const int& dim, const int& nst)->void { - if (GlobalV::MY_RANK == 0) { assert(nst == LR_Util::read_value(efile(label), e, nst)); } - assert(nst * dim == LR_Util::read_value(vfile(label), v, nst, dim)); + if (GlobalV::MY_RANK == 0) { + assert(nst == LR_Util::read_value(efile_in(label), e, nst)); + std::cout <<"Rank "<< GlobalV::MY_RANK << ": finish reading " << efile_in(label) << std::endl; + } +#ifdef __MPI +// in velocity gauge, the eigenvalues may be used to calculate the transition dipole, so we'd better broadcast them + MPI_Bcast(e, nst, MPI_DOUBLE, 0, MPI_COMM_WORLD); +#endif + assert(nst * dim == LR_Util::read_value(vfile_in(label), v, nst, dim)); + std::cout <<"Rank "<< GlobalV::MY_RANK << ": finish reading " << vfile_in(label) << std::endl; }; - std::cout << "reading the excitation amplitudes from file: \n"; + std::cout << "reading the excitation states from file: \n"; if (openshell) { - read_states("openshell", this->pelec->ekb.c, this->X[0].template data(), nloc_per_band, nstates); + read_states("openshell", this->pelec->ekb.c, this->X[0].template data(), nloc_per_state, nstates); } else { auto spin_types = std::vector({ "singlet", "triplet" }); - for (int is = 0;is < nspin;++is) { read_states(spin_types[is], this->pelec->ekb.c + is * nstates, this->X[is].template data(), nloc_per_band, nstates); } + for (int is = 0;is < nspin;++is) { read_states(spin_types[is], this->pelec->ekb.c + is * nstates, this->X[is].template data(), nloc_per_state, nstates); } } } ModuleBase::timer::end("ESolver_LR", "runner"); @@ -566,6 +618,13 @@ void LR::ESolver_LR::after_all_runners(UnitCell& ucell) ModuleBase::TITLE("ESolver_LR", "after_all_runners"); if (input.ri_hartree_benchmark != "none") { return; } //no need to calculate the spectrum in the benchmark routine //cal spectrum + if (LR_Util::tolower(this->input.abs_gauge) == "velocity" ) + { + const int nspin_tmp = PARAM.inp.nspin == 2 ? 2 : 1; + this->velocity_mo = LR_Util::cal_velocity_mo(*this->ucell_, this->gd, this->two_center_bundle_, + this->paraMat_, this->paraC_, this->kv, *this->psi_ks, + this->nk, nspin_tmp, this->nbasis, this->nocc, this->nvirt); + } std::vector freq(100); std::vector abs_wavelen_range({ 20, 200 });//default range if (input.abs_wavelen_range.size() >= 2 && std::abs(input.abs_wavelen_range[1] - input.abs_wavelen_range[0]) > 0.02) @@ -581,16 +640,24 @@ void LR::ESolver_LR::after_all_runners(UnitCell& ucell) LR_Spectrum spectrum(nspin, this->nbasis, this->nocc, this->nvirt, *this->pw_rho, *this->psi_ks, *this->ucell_, this->kv, this->gd, this->orb_cutoff_, this->two_center_bundle_, this->paraX_, this->paraC_, this->paraMat_, - &this->pelec->ekb.c[is * nstates], this->X[is].template data(), nstates, openshell, + &this->pelec->ekb.c[is * nstates], this->eig_ks.c, this->X[is].template data(), nstates, openshell, LR_Util::tolower(input.abs_gauge)); - spectrum.transition_analysis(spin_types[is]); + if (LR_Util::tolower(this->input.abs_gauge) == "velocity" ) {spectrum.set_vmo(this->velocity_mo.data());} + spectrum.cal_spectrum(); + spectrum.transition_analysis(spin_types[is]+"_tda"); if (spin_types[is] != "triplet") // triplets has no transition dipole and no contribution to the spectrum { spectrum.optical_absorption_method1(freq, input.abs_broadening); + spectrum.write_transition_dipole(PARAM.globalv.global_out_dir + + "trans_dipole_" + spin_types[is] + "_tda.dat"); // =============================================== for test ==================================================== // spectrum.optical_absorption_method2(freq, input.abs_broadening); - // spectrum.test_transition_dipoles_velocity_ks(eig_ks.c); - // spectrum.write_transition_dipole(PARAM.globalv.global_out_dir + "dipole_velocity_ks.dat"); + // if (LR_Util::tolower(input.abs_gauge) == "velocity") + // { // TEST the formula v/omega rather than v/(e_a-e_i) + // spectrum.test_transition_dipoles_velocity_omega(); + // spectrum.write_transition_dipole(PARAM.globalv.global_out_dir + + // "trans_dipole_" + spin_types[is] + "_vomega_tda.dat"); + // } // =============================================== for test ==================================================== } } @@ -610,9 +677,9 @@ void LR::ESolver_LR::setup_eigenvectors_X() );//nvirt - row, nocc - col this->paraX_.emplace_back(std::move(px)); } - this->nloc_per_band = nk * (openshell ? paraX_[0].get_local_size() + paraX_[1].get_local_size() : paraX_[0].get_local_size()); + this->nloc_per_state = nk * (openshell ? paraX_[0].get_local_size() + paraX_[1].get_local_size() : paraX_[0].get_local_size()); - this->X.resize(openshell ? 1 : nspin, LR_Util::newTensor({ nstates, nloc_per_band })); + this->X.resize(openshell ? 1 : nspin, LR_Util::newTensor({ nstates, nloc_per_state })); for (auto& x : X) { x.zero(); } auto spin_types = (nspin == 2 && !openshell) ? std::vector({ "singlet", "triplet" }) : std::vector({ "updown" }); @@ -654,7 +721,7 @@ void LR::ESolver_LR::set_X_initial_guess() const int occ_global = std::get<0>(ix2ioiv[ipair]); // occ const int virt_global = std::get<1>(ix2ioiv[ipair]); // virt const int ik = ib / np; - const int xstart_b = ib * nloc_per_band; //start index of band ib + const int xstart_b = ib * nloc_per_state; //start index of band ib const int xstart_bs = (openshell && is == 1) ? xstart_b + nk * paraX_[0].get_local_size() : xstart_b; // start index of band ib, spin is const int is_in_x = openshell ? 0 : is; // if openshell, spin-up and spin-down are put together if (px.in_this_processor(virt_global, occ_global)) @@ -711,8 +778,9 @@ void LR::ESolver_LR::read_ks_wfc() else if (!ModuleIO::read_wfc_nao(PARAM.globalv.global_readin_dir, this->paraMat_, *this->psi_ks, this->pelec->ekb, this->pelec->wg, - this->pelec->klist->ik2iktot, - this->pelec->klist->get_nkstot(), + this->kv.ik2iktot, + this->kv.get_nkstot(), + PARAM.inp.nspin, /*skip_bands=*/this->nocc_max - this->nocc_in)) { ModuleBase::WARNING_QUIT("ESolver_LR", "read ground-state wavefunction failed."); } @@ -729,9 +797,8 @@ void LR::ESolver_LR::read_ks_chg(Charge& chg_gs) for (int is = 0; is < this->nspin; ++is) { std::stringstream ssc; - ssc << PARAM.globalv.global_readin_dir << "SPIN" << is + 1 << "_CHG.cube"; + ssc << PARAM.globalv.global_readin_dir << "chgs" << is + 1 << ".cube"; GlobalV::ofs_running << ssc.str() << std::endl; - double ef = 0.0; if (ModuleIO::read_vdata_palgrid(Pgrid, GlobalV::MY_RANK, GlobalV::ofs_running, @@ -742,9 +809,9 @@ void LR::ESolver_LR::read_ks_chg(Charge& chg_gs) } else { // prenspin for nspin=4 is not supported currently ModuleBase::WARNING_QUIT( "init_rho", - "!!! Couldn't find the charge file !!! The default directory \n of SPIN1_CHG.cube is OUT.suffix, " + "!!! Couldn't find the charge file !!! The default directory \n of " + ssc.str() +" is OUT.suffix, " "or you must set read_file_dir \n to a specific directory. "); -} + } } } template class LR::ESolver_LR; diff --git a/source/source_lcao/module_lr/esolver_lrtd_lcao.h b/source/source_lcao/module_lr/esolver_lrtd_lcao.h index 92c0b45fcd..1029f90802 100644 --- a/source/source_lcao/module_lr/esolver_lrtd_lcao.h +++ b/source/source_lcao/module_lr/esolver_lrtd_lcao.h @@ -61,11 +61,11 @@ namespace LR ModuleBase::matrix eig_ks;///< energy of ground state /// @brief Excited state wavefunction (locc, lvirt are local size of nocc and nvirt in each process) - /// size of X: [neq][{nstate, nloc_per_band}], namely: + /// size of X: [neq][{nstate, nloc_per_state}], namely: /// - [nspin][{nstates, nk* (locc* lvirt}] for close- shell, /// - [1][{nstates, nk * (locc[0] * lvirt[0]) + nk * (locc[1] * lvirt[1])}] for open-shell std::vector X; - int nloc_per_band = 1; + int nloc_per_state = 1; std::vector nocc; ///< number of occupied orbitals for each spin used in the calculation int nocc_in = 1; ///< nocc read from input (adjusted by nelec): max(spin-up, spindown) @@ -101,16 +101,21 @@ namespace LR TwoCenterBundle two_center_bundle_; + LCAO_Orbitals orb_; ///< numerical atomic orbital data for single-point evaluation + std::vector> velocity_mo; ///< store the velocity matrix elements in MO basis + int cal_nupdown_form_occ(const ModuleBase::matrix& wg); + void setup_2center_table(TwoCenterBundle& two_center_bundle, LCAO_Orbitals& orb, UnitCell& ucell); + /// @brief allocate and set the inital value of X void setup_eigenvectors_X(); void set_X_initial_guess(); /// @brief read in the ground state wave function, band energy and occupation - void read_ks_wfc(); + virtual void read_ks_wfc(); /// @brief read in the ground state charge density void read_ks_chg(Charge& chg); - void init_pot(const Charge& chg_gs); + virtual void init_pot(const Charge& chg_gs); /// @brief check the legality of the input parameters void parameter_check() const; diff --git a/source/source_lcao/module_lr/hamilt_casida.h b/source/source_lcao/module_lr/hamilt_casida.h index b9459cfa8a..61d546602b 100644 --- a/source/source_lcao/module_lr/hamilt_casida.h +++ b/source/source_lcao/module_lr/hamilt_casida.h @@ -10,7 +10,7 @@ #include "source_lcao/module_lr/operator_casida/operator_lr_exx.h" #include "source_lcao/module_lr/ri_benchmark/operator_ri_hartree.h" #include "source_lcao/module_ri/LRI_CV_Tools.h" -#include "source_lcao/module_lr/ri_benchmark/ri_benchmark.h" +#include "source_lcao/module_lr/utils/lr_io.h" #endif namespace LR { diff --git a/source/source_lcao/module_lr/utils/exciton_plotter.h b/source/source_lcao/module_lr/utils/exciton_plotter.h index f49126820c..9046e81f8c 100644 --- a/source/source_lcao/module_lr/utils/exciton_plotter.h +++ b/source/source_lcao/module_lr/utils/exciton_plotter.h @@ -6,7 +6,7 @@ #include "source_cell/klist.h" #include "source_estate/module_dm/density_matrix.h" #include "source_io/module_output/cube_io.h" -#include "source_lcao/module_gint/gint_interface.h" +#include "source_hamilt/module_gint/gint_interface.h" #include "source_lcao/module_lr/dm_trans/dm_trans.h" #include "source_lcao/module_lr/utils/lr_util.h" #include "source_lcao/module_lr/utils/lr_util_hcontainer.h" diff --git a/source/source_lcao/module_lr/utils/spectrum_mo.hpp b/source/source_lcao/module_lr/utils/spectrum_mo.hpp index d0b3e37275..8e7b70e53d 100644 --- a/source/source_lcao/module_lr/utils/spectrum_mo.hpp +++ b/source/source_lcao/module_lr/utils/spectrum_mo.hpp @@ -3,7 +3,7 @@ #include "source_basis/module_nao/two_center_bundle.h" #include "source_cell/klist.h" #include "source_io/module_parameter/parameter.h" -#include "source_lcao/module_hcontainer/hcontainer_funcs.h" +#include "source_hamilt/module_hcontainer/hcontainer_funcs.h" #include "source_lcao/module_lr/ao_to_mo_transformer/ao_to_mo.h" #include "source_lcao/module_lr/utils/lr_util.h" #include "source_lcao/module_lr/utils/lr_util_print.h"