diff --git a/include/RI/distribute/Distribute_Equally.h b/include/RI/distribute/Distribute_Equally.h index 3f1acca..2e58229 100644 --- a/include/RI/distribute/Distribute_Equally.h +++ b/include/RI/distribute/Distribute_Equally.h @@ -56,6 +56,27 @@ namespace Distribute_Equally const std::array &period, const std::size_t num_index, const bool flag_task_repeatable); + + template + extern void distribute_atom_and_k_pair( + const MPI_Comm &mpi_comm, + const std::size_t nat, + const std::size_t nk, + std::vector &list_I, + std::vector &list_J, + std::vector &list_k1_index, + std::vector &list_k2_index, + const bool flag_task_repeatable); + + template + extern void distribute_atom_pair_and_k( + const MPI_Comm &mpi_comm, + const std::size_t nat, + const std::size_t nk, + std::vector &list_I, + std::vector &list_J, + std::vector &list_k_index, + const bool flag_task_repeatable); } } diff --git a/include/RI/distribute/Distribute_Equally.hpp b/include/RI/distribute/Distribute_Equally.hpp index 3dbdaf8..e987884 100644 --- a/include/RI/distribute/Distribute_Equally.hpp +++ b/include/RI/distribute/Distribute_Equally.hpp @@ -126,6 +126,136 @@ namespace Distribute_Equally period); return atoms_split_list; } + + // 均分{atomI,atomJ,k1,k2} + template + void distribute_atom_and_k_pair( + const MPI_Comm &mpi_comm, + const std::size_t nat, + const std::size_t nk, + std::vector &list_I, + std::vector &list_J, + std::vector &list_k1_index, + std::vector &list_k2_index, + const bool flag_task_repeatable) + { + // task_sizes的顺序必须从小到大，否则在split中会出现rank_size *A1_ptr, *A2_ptr, *B1_ptr, *B2_ptr; + if (nk >= nat) + { + ntaskA = nat; + ntaskB = nk; + A1_ptr = &list_I; + A2_ptr = &list_J; + B1_ptr = &list_k1_index; + B2_ptr = &list_k2_index; + } + else + { + ntaskA = nk; + ntaskB = nat; + A1_ptr = &list_k1_index; + A2_ptr = &list_k2_index; + B1_ptr = &list_I; + B2_ptr = &list_J; + } + const std::vector task_sizes{ntaskA, ntaskA, ntaskB, ntaskB}; + const std::vector> + comm_color_sizes = Split_Processes::split_all(mpi_comm, task_sizes); + + if(!flag_task_repeatable) + if(RI::MPI_Wrapper::mpi_get_rank(std::get<0>(comm_color_sizes.back())())) + return; + + std::vector indicesA, indicesB; + for(Tindex i=0; i(comm_color_sizes[1]), + std::get<2>(comm_color_sizes[1]), + indicesA); + *A2_ptr = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[2]), + std::get<2>(comm_color_sizes[2]), + indicesA); + *B1_ptr = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[3]), + std::get<2>(comm_color_sizes[3]), + indicesB); + *B2_ptr = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[4]), + std::get<2>(comm_color_sizes[4]), + indicesB); + } + + // 均分{atomI,atomJ,k} + template + void distribute_atom_pair_and_k( + const MPI_Comm &mpi_comm, + const std::size_t nat, + const std::size_t nk, + std::vector &list_I, + std::vector &list_J, + std::vector &list_k_index, + const bool flag_task_repeatable) + { + std::vector task_sizes; + std::vector indices_atom, indices_k; + for(Tindex i=0; i= nat) + { + task_sizes = {nat, nat, nk}; + } + else + { + task_sizes = {nk, nat, nat}; + } + const std::vector> + comm_color_sizes = Split_Processes::split_all(mpi_comm, task_sizes); + + if(!flag_task_repeatable) + if(RI::MPI_Wrapper::mpi_get_rank(std::get<0>(comm_color_sizes.back())())) + return; + + if (nk >= nat) + { + list_I = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[1]), + std::get<2>(comm_color_sizes[1]), + indices_atom); + list_J = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[2]), + std::get<2>(comm_color_sizes[2]), + indices_atom); + list_k_index = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[3]), + std::get<2>(comm_color_sizes[3]), + indices_k); + } + else + { + list_k_index = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[1]), + std::get<2>(comm_color_sizes[1]), + indices_k); + list_I = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[2]), + std::get<2>(comm_color_sizes[2]), + indices_atom); + list_J = Divide_Atoms::divide_atoms( + std::get<1>(comm_color_sizes[3]), + std::get<2>(comm_color_sizes[3]), + indices_atom); + } + } } } \ No newline at end of file diff --git a/include/RI/global/Array_Operator.h b/include/RI/global/Array_Operator.h index 94132ad..eea17df 100644 --- a/include/RI/global/Array_Operator.h +++ b/include/RI/global/Array_Operator.h @@ -13,7 +13,12 @@ namespace RI namespace Array_Operator { template - extern std::array operator%(const std::array &v1, const std::array &v2); + typename std::enable_if::value, std::array>::type + operator%(const std::array &v1, const std::array &v2); + + template + typename std::enable_if::value, std::array>::type + operator%(const std::array &v1, const std::array &v2); template extern std::array operator+(const std::array &v1, const std::array &v2); diff --git a/include/RI/global/Array_Operator.hpp b/include/RI/global/Array_Operator.hpp index 4e87f83..18d372a 100644 --- a/include/RI/global/Array_Operator.hpp +++ b/include/RI/global/Array_Operator.hpp @@ -13,7 +13,8 @@ namespace RI namespace Array_Operator { template - std::array operator%(const std::array &v1, const std::array &v2) + typename std::enable_if::value, std::array>::type + operator%(const std::array &v1, const std::array &v2) { auto mod = [](const T i, const T n){ return (i%n+3*n/2)%n-n/2; }; // [-n/2,n/2] // auto mod = [](const T i, const T n){ return (i%n+n)%n; }; // [0,n) @@ -23,6 +24,24 @@ namespace Array_Operator v[i] = mod(v1[i], v2[i]); return v; } + template + typename std::enable_if::value, std::array>::type + operator%(const std::array &v1, const std::array &v2) + { + constexpr double epsilon = 1e-6; + auto mod_f = [&](T x, T period){ // [0, period) + T r = std::fmod(x, period); + if (std::abs(r) < epsilon) r = 0.0; + if (std::abs(r - period) < epsilon) r = 0.0; + if (r < 0) r += period; + + return r; + }; + std::array v; + for(std::size_t i=0; i std::array operator+(const std::array &v1, const std::array &v2) diff --git a/include/RI/global/Global_Func-1.h b/include/RI/global/Global_Func-1.h index ce23b0f..2dcf699 100644 --- a/include/RI/global/Global_Func-1.h +++ b/include/RI/global/Global_Func-1.h @@ -69,6 +69,13 @@ namespace Global_Func else return ptr->second; } + + template + const T &find(const T &data) + { + return data; + } + template inline const auto &find( const std::map> &m, @@ -121,6 +128,21 @@ namespace Global_Func { return std::vector(v.begin(), v.end()); } + + /// @brief sorted unique union of two containers + /// @tparam C1,C2 containers supporting .begin()/.end() and iterator-pair construction + /// (e.g., std::vector, std::deque, std::list) + template + static auto set_union(const C1& c1, const C2& c2) + -> C1 + { + static_assert( + std::is_same::value, + "set_union: both containers must have the same value_type"); + std::set s(c1.begin(), c1.end()); + s.insert(c2.begin(), c2.end()); + return C1(s.begin(), s.end()); + } } } \ No newline at end of file diff --git a/include/RI/global/Global_Func-2.h b/include/RI/global/Global_Func-2.h index ade8c68..30da735 100644 --- a/include/RI/global/Global_Func-2.h +++ b/include/RI/global/Global_Func-2.h @@ -49,6 +49,21 @@ namespace Global_Func typename std::enable_if::value,int>::type =0> Tout convert(const Tin &t) { return t.real(); } + + template< + typename T, + typename std::enable_if::value,int>::type =0> + inline T get_conj(const T& x) + { + return x; + } + template< + typename T, + typename std::enable_if< Global_Func::is_complex::value,int>::type =0> + inline T get_conj(const T& x) + { + return std::conj(x); + } } } \ No newline at end of file diff --git a/include/RI/global/Shape_Vector.h b/include/RI/global/Shape_Vector.h index 7e97691..ed6535b 100644 --- a/include/RI/global/Shape_Vector.h +++ b/include/RI/global/Shape_Vector.h @@ -29,7 +29,12 @@ class Shape_Vector for(auto ptr_in=v_in.begin(); ptr_in& v_in) + :size_(v_in.size()) + { + assert(v_in.size() <= sizeof(v) / sizeof(*v)); + for (std::size_t i = 0;i < size_;++i) this->v[i] = v_in[i]; + } const std::size_t* begin() const noexcept { return this->v; } const std::size_t* end() const noexcept { return this->v+size_; } std::size_t size() const noexcept { return size_; } diff --git a/include/RI/global/Tensor.h b/include/RI/global/Tensor.h index 4f5a2c9..dea7ef5 100644 --- a/include/RI/global/Tensor.h +++ b/include/RI/global/Tensor.h @@ -45,6 +45,7 @@ class Tensor Tensor transpose() const; Tensor dagger() const; + Tensor conjugate() const; // ||d||_p = (|d_1|^p+|d_2|^p+...)^{1/p} // if(p==std::numeric_limits::max()) ||d||_max = max_i |d_i| diff --git a/include/RI/global/Tensor.hpp b/include/RI/global/Tensor.hpp index dec3930..6c55ced 100644 --- a/include/RI/global/Tensor.hpp +++ b/include/RI/global/Tensor.hpp @@ -13,6 +13,7 @@ #include #include #include +#include #include namespace RI @@ -95,10 +96,18 @@ template bool same_shape (const Tensor &t1, const Tensor &t2) { if(t1.shape.size() != t2.shape.size()) + { + std::cerr << "same_shape: ndim mismatch (" + << t1.shape.size() << " vs " << t2.shape.size() << ")" << std::endl; return false; + } for(std::size_t ishape=0; ishape Tensor::dagger() const return t; } +template +Tensor Tensor::conjugate() const +{ + Tensor t(this->shape); + for (std::size_t i = 0; i < this->data->size(); ++i) + (*t.data)[i] = Global_Func::get_conj((*this->data)[i]); + return t; +} + template Global_Func::To_Real_t Tensor::norm(const double p) const { diff --git a/include/RI/global/Tensor_Multiply-gemv.hpp b/include/RI/global/Tensor_Multiply-gemv.hpp new file mode 100644 index 0000000..bdde598 --- /dev/null +++ b/include/RI/global/Tensor_Multiply-gemv.hpp @@ -0,0 +1,52 @@ +#include "Tensor_Multiply.h" + +namespace RI +{ + +namespace Tensor_Multiply +{ + // Txy(x0,x1,..., xM) = Tx(x0, x1,x2, ..., xM, y0, y1, ..., yN) * Vy(y0, y1, ..., yN) + template + Tensor gemv(const Tensor& Tx, const Tensor& Vy) + { + assert(Tx.shape.size() >= Vy.shape.size()); + const std::size_t ny = Vy.get_shape_all(); + assert(Tx.get_shape_all() % ny == 0); + const std::size_t dim = Tx.shape.size() - Vy.shape.size(); + std::vector shape_vector; + if (dim == 0) + shape_vector.push_back(1); + else + for (int d = 0; d < dim; ++d) + shape_vector.push_back(Tx.shape[d]); + Tensor Txy(shape_vector); + Blas_Interface::gemv( + 'N', Txy.get_shape_all(), ny, + Tdata(1.0), Tx.ptr(), Vy.ptr(), + Tdata(0.0), Txy.ptr()); + return Txy; + } + + // Txy(x0,x1,..., xM) = Tx(y0, y1,y2, ..., yN, x0, x1, ..., xM) * Vy(y0, y1, ..., yN) + template + Tensor gemv_trans(const Tensor& Tx, const Tensor& Vy) + { + assert(Tx.shape.size() >= Vy.shape.size()); + const std::size_t ny = Vy.get_shape_all(); + assert(Tx.get_shape_all() % ny == 0); + const std::size_t dim = Tx.shape.size() - Vy.shape.size(); + std::vector shape_vector; + if (dim == 0) + shape_vector.push_back(1); + else + for (int d = 0; d < dim; ++d) + shape_vector.push_back(Tx.shape[Vy.shape.size() + d]); + Tensor Txy(shape_vector); + Blas_Interface::gemv( + 'T', ny, Txy.get_shape_all(), + Tdata(1.0), Tx.ptr(), Vy.ptr(), + Tdata(0.0), Txy.ptr()); + return Txy; + } +} +} \ No newline at end of file diff --git a/include/RI/global/Tensor_Multiply.h b/include/RI/global/Tensor_Multiply.h index 4d7789e..e762202 100644 --- a/include/RI/global/Tensor_Multiply.h +++ b/include/RI/global/Tensor_Multiply.h @@ -11,4 +11,5 @@ #include "Tensor_Multiply-22.hpp" #include "Tensor_Multiply-23.hpp" #include "Tensor_Multiply-32.hpp" -#include "Tensor_Multiply-33.hpp" \ No newline at end of file +#include "Tensor_Multiply-33.hpp" +#include "Tensor_Multiply-gemv.hpp" \ No newline at end of file diff --git a/include/RI/physics/Hartree.h b/include/RI/physics/Hartree.h new file mode 100644 index 0000000..e147471 --- /dev/null +++ b/include/RI/physics/Hartree.h @@ -0,0 +1,121 @@ +// =================== +// Author: Ziqing Guan +// date: 2025.12.26 +// =================== + +#pragma once +#include "../global/Global_Func-2.h" +#include "../global/Tensor.h" +#include "../global/Global_Func-1.h" +#include "../ri/LRI_k.h" + +#include +#include +#include +#include + +namespace RI +{ +// feat: calculate hartree term for qs-GW based on cvcd in k space +template +class Hartree +{ +public: + using TC = std::array; + using TAC = std::pair; + using Tdata_real = Global_Func::To_Real_t; + using Tk = std::array; + + Hartree () = default; + + void init(std::vector kindex_map_in) + { + this->kindex_map = std::move(kindex_map_in); + } + + void set_parallel( + const MPI_Comm &mpi_comm_in, const std::size_t nat, const std::size_t nk, + const std::array &period_in) + { + this->lrik.mpi_comm = mpi_comm_in; + this->lrik.period = period_in; + RI::Distribute_Equally::distribute_atom_pair_and_k(mpi_comm_in, + nat, nk, this->list_I, this->list_J, this->k_indices, false); + + this->list_IJ = Global_Func::set_union(this->list_I, this->list_J); + this->flag_finish.stru = true; + } + + void set_Cs( + std::map>> &Cs, + const Tdata_real &threshold, + const std::set &listI, + const std::set &listJ, + const std::string &save_name_suffix="") + { + Cs = Communicate_Tensors_Map_Judge::comm_map2_first(this->lrik.mpi_comm, std::move(Cs), listI, listJ); + this->lrik.set_tensors_map2( + Cs, + {Label::ab::a, Label::ab::b}, + {{"threshold_filter", threshold}, {"flag_comm", false}}, + "Cs_"+save_name_suffix ); + this->flag_finish.Cs = true; + } + + void free_Cs(const std::string &save_name_suffix="") + { + this->lrik.free_tensors_map2("Cs_"+save_name_suffix); + this->flag_finish.Cs = false; + }; + + void set_Vs( + std::map>> &Vs, + const Tdata_real &threshold, + const std::set &listI, + const std::set &listJ, + const std::string &save_name_suffix="") + { + Vs = Communicate_Tensors_Map_Judge::comm_map2_first(this->lrik.mpi_comm, std::move(Vs), listI, listJ); + this->lrik.set_tensors_map2( + Vs, + {Label::ab::a0b0}, + {{"threshold_filter", threshold}, {"flag_comm", false}}, + "Vs_"+save_name_suffix ); + this->flag_finish.Vs = true; + }; + + void free_Vs(const std::string &save_name_suffix="") + { + this->lrik.free_tensors_map2("Vs_"+save_name_suffix); + this->flag_finish.Vs = false; + }; + + std::map>>> cal_hartree( + const std::map>>>& Ds, + const std::string &save_name_C="Cs_", const std::string &save_name_V="Vs_") + { + return this->lrik.cal_cvcd_k_hartree( + Ds, this->kindex_map, this->k_indices, this->list_I, this->list_J, this->list_IJ, + save_name_C, save_name_V); + } + + std::vector kindex_map; // index → Tk fractional coord + std::vector k_indices; + std::vector list_I; + std::vector list_J; + std::vector list_IJ; // I∪J + + LRI_k lrik; + +private: + struct Flag_Finish + { + bool stru=false; + bool Cs=false; + bool Vs=false; + }; + Flag_Finish flag_finish; + +}; + +} diff --git a/include/RI/physics/LR.h b/include/RI/physics/LR.h new file mode 100644 index 0000000..373e815 --- /dev/null +++ b/include/RI/physics/LR.h @@ -0,0 +1,181 @@ +#pragma once +#include "../global/Global_Func-2.h" +#include "../global/Tensor.h" +#include "../global/Global_Func-1.h" +#include "../ri/LRI_k.h" + +#include +#include +#include +#include +namespace RI +{ +template +class LR +{ +public: + using TC = std::array; + using TAC = std::pair; + using Tdata_real = Global_Func::To_Real_t; + using Tk = std::array; + + LR() = default; + + void init(std::vector kindex_map_in, int nocc_in, int nvirt_in) + { + this->kindex_map = std::move(kindex_map_in); + this->nocc = nocc_in; + this->nvirt = nvirt_in; + } + + void set_parallel( + const MPI_Comm &mpi_comm_in, const std::size_t nat, const std::size_t nk, + const std::array &period_in) + { + this->lrik.mpi_comm = mpi_comm_in; + this->lrik.period = period_in; + + RI::Distribute_Equally::distribute_atom_and_k_pair(mpi_comm_in, + nat, nk, this->list_I, this->list_J, + this->k1_indices, this->k2_indices, false); + + this->list_IJ = Global_Func::set_union(this->list_I, this->list_J); + this->k_indices = Global_Func::set_union(this->k1_indices, this->k2_indices); + + this->flag_finish.stru = true; + } + + void set_Cs( + std::map>> &Cs, + const Tdata_real &threshold, + const std::set &listIJ, + const std::set &all_atoms, + const std::string &save_name_suffix="") + { + //Cs = Communicate_Tensors_Map_Judge::comm_map2_first(this->lrik.mpi_comm, std::move(Cs), listI, listJ); + this->lrik.set_tensors_map2( + Cs, + {Label::ab::a, Label::ab::b}, + {{"threshold_filter", threshold}, {"flag_comm", false}}, + "Cs_"+save_name_suffix ); + this->flag_finish.Cs = true; + } + void free_Cs(const std::string &save_name_suffix="") + { + this->lrik.free_tensors_map2("Cs_"+save_name_suffix); + this->flag_finish.Cs = false; + }; + + void set_Vs( + std::map>> &Vs, + const Tdata_real &threshold, + const std::set &listI, + const std::set &listJ, + const std::string &save_name_suffix="") + { + //Vs = Communicate_Tensors_Map_Judge::comm_map2_first(this->lrik.mpi_comm, std::move(Vs), listI, listJ); + this->lrik.set_tensors_map2( + Vs, + {Label::ab::a0b0}, + {{"threshold_filter", threshold}, {"flag_comm", false}}, + "Vs_"+save_name_suffix ); + this->flag_finish.Vs = true; + }; + void free_Vs(const std::string &save_name_suffix="") + { + this->lrik.free_tensors_map2("Vs_"+save_name_suffix); + this->flag_finish.Vs = false; + }; + + void set_Ws( + std::map>> &Ws, + const Tdata_real &threshold, + const std::set &listI, + const std::set &listJ, + const std::string &save_name_suffix="") + { + //Ws = Communicate_Tensors_Map_Judge::comm_map2_first(this->lrik.mpi_comm, std::move(Ws), listI, listJ); + this->lrik.set_tensors_map2( + Ws, + {Label::ab::a0b0}, + {{"threshold_filter", threshold}, {"flag_comm", false}}, + "Ws_"+save_name_suffix ); + this->flag_finish.Ws = true; + }; + void free_Ws(const std::string &save_name_suffix="") + { + this->lrik.free_tensors_map2("Ws_"+save_name_suffix); + this->flag_finish.Ws = false; + }; + + /// @brief calculate Csk_ao_mo on-the-fly and store internally + void cal_Csk_ao_mo( + const std::string& save_name, + std::ofstream& ofs) + { + const auto& CsR_ao = this->lrik.data_pool.at(save_name).Ds_ab; + this->Csk_ao_mo = this->lrik.cal_Csk_ao_mo( + CsR_ao, this->map_psi, this->kindex_map, + this->k_indices, this->list_IJ, ofs); + } + + std::map>> + cal_cvc_mo_k_onthefly( + const std::vector& psi_type, + const std::string& save_name, + const bool is_A) + { + return this->lrik.cal_cvc_mo_k_onthefly( + this->Csk_ao_mo, this->map_psi, + this->k1_indices, this->k2_indices, + this->list_I, this->list_J, + psi_type, this->nocc, this->nvirt, + save_name, is_A, + this->q_list, this->q2kpair); + } + + std::map>> + cal_cvc_mo_k_hartree_onthefly( + const std::vector& psi_type, + const std::string& save_name, + const bool is_A) + { + return this->lrik.cal_cvc_mo_k_hartree_onthefly( + this->Csk_ao_mo, this->map_psi, + this->k1_indices, this->k2_indices, + this->list_I, this->list_J, + psi_type, this->nocc, this->nvirt, + save_name, is_A); + } + + std::vector kindex_map; // index → Tk fractional coord + std::size_t nocc; + std::size_t nvirt; + + std::vector k1_indices; + std::vector k2_indices; + std::vector k_indices; // k1∪k2 (for psi/mo transform) + std::vector list_I; + std::vector list_J; + std::vector list_IJ; // I∪J + + std::vector q_list; + std::map>> q2kpair; // q → list of (k1_idx, k2_idx) + + std::map>> map_psi; //> + std::map>> Csk_ao_mo; + LRI_k lrik; + +private: + struct Flag_Finish + { + bool stru=false; + bool Cs=false; + bool Vs=false; + bool Ws=false; + }; + Flag_Finish flag_finish; + +}; + +} diff --git a/include/RI/ri/Cell_Nearest.h b/include/RI/ri/Cell_Nearest.h index a317346..66b67e9 100644 --- a/include/RI/ri/Cell_Nearest.h +++ b/include/RI/ri/Cell_Nearest.h @@ -26,6 +26,10 @@ class Cell_Nearest TC get_cell_nearest_discrete(const TA &Ax, const TA &Ay, const TC &cell) const; + /// @brief get nearest cell through brute-force search + /// @details if more than one cell has the min distance, choose the direction as smaller abs(R-R_near).z, then .y, then .x + TC cell_nearest_direction(const TA Ax, const TA Ay, const TC& cell, double& dist_min) const; + public: //private: TC period; std::map> atoms_pos; diff --git a/include/RI/ri/Cell_Nearest.hpp b/include/RI/ri/Cell_Nearest.hpp index cb6a9ef..8602405 100644 --- a/include/RI/ri/Cell_Nearest.hpp +++ b/include/RI/ri/Cell_Nearest.hpp @@ -79,4 +79,57 @@ auto Cell_Nearest::get_cell_nearest_discrete( return cell_nearest + cell; } +template +auto Cell_Nearest::cell_nearest_direction( + const TA Ax, const TA Ay, const TC &cell, double &dist_min) const +-> TC +{ + static_assert(Ndim == 3, "cell_nearest_direction currently assumes Ndim==3."); + + TC cell_nearest = cell; + const std::array &Ryx = this->cells_nearest_continuous.at(Ax).at(Ay); //frac coordinate, Rx-Ry + TC cell_try; + Tensor diff({Ndim}); // frac coordinate, pos_y - pos_x + + for (std::size_t i = 0; i < Ndim; ++i) + diff(i) = cell_nearest[i] - Ryx[i]; + int a_min(0), b_min(0), c_min(0); + dist_min = (diff * this->latvec).norm(2); + + for (int a = -2; a < 3; ++a) + { + cell_try[0] = a * this->period[0] + cell[0]; + diff(0) = cell_try[0] - Ryx[0]; + for (int b = -2; b < 3; ++b) + { + cell_try[1] = b * this->period[1] + cell[1]; + diff(1) = cell_try[1] - Ryx[1]; + for (int c = -2; c < 3; ++c) + { + cell_try[2] = c * this->period[2] + cell[2]; + diff(2) = cell_try[2] - Ryx[2]; + double dist = (diff * this->latvec).norm(2); + if (dist < dist_min - 1e-6) + { + dist_min = dist; + cell_nearest = cell_try; + a_min = a; b_min = b; c_min = c; + } + else if (std::abs(dist - dist_min) < 1e-6) + { + // in case of tie, choose the one with smaller abs(R-R_near).z, then .y, then .x + if (std::abs(c) < std::abs(c_min) || + (std::abs(c) == std::abs(c_min) && (std::abs(b) < std::abs(b_min) || + (std::abs(b) == std::abs(b_min) && std::abs(a) < std::abs(a_min) )))) + { + dist_min = dist; + cell_nearest = cell_try; + a_min = a; b_min = b; c_min = c; + } + } + } + } + } + return cell_nearest; +} } diff --git a/include/RI/ri/LRI.h b/include/RI/ri/LRI.h index 5cf0c1b..6e9c3e5 100644 --- a/include/RI/ri/LRI.h +++ b/include/RI/ri/LRI.h @@ -93,4 +93,4 @@ class LRI #include "LRI.hpp" #include "LRI-set.hpp" -#include "LRI-cal_loop3.hpp" +#include "LRI-cal_loop3.hpp" \ No newline at end of file diff --git a/include/RI/ri/LRI_Cal_Aux.h b/include/RI/ri/LRI_Cal_Aux.h index 1708ed1..a5dcb4f 100644 --- a/include/RI/ri/LRI_Cal_Aux.h +++ b/include/RI/ri/LRI_Cal_Aux.h @@ -83,6 +83,24 @@ namespace LRI_Cal_Aux } } + // D_result = D_add * scale (init) + // or D_result += D_add * scale (accumulate via BLAS axpy) + template + inline void add_Ds( + const Tensor &D_add, + Tensor &D_result, + const Tdata scale) + { + if (D_result.empty()) + { + D_result = D_add * scale; + } + else + { + Blas_Interface::axpy(scale, D_add, D_result); + } + } + template void add_Ds( std::map &&Ds_add, @@ -160,16 +178,18 @@ namespace LRI_Cal_Aux } } - template + // Tkey labels thread lock, in actual use, Tkey can be TA, TC, Tk, etc. + // Tvalue can be Tensor or another map, both of them can be input of add_Ds + template void add_Ds_omp_try_map( - std::map>> &Ds_result_thread, - std::map>> &Ds_result, - std::map &lock_Ds_result_add_map, + std::map& Ds_result_thread, + std::map& Ds_result, + std::map &lock_Ds_result_add_map, const double &fac) { for(auto ptr=Ds_result_thread.begin(); ptr!=Ds_result_thread.end(); ) { - const TA key = ptr->first; + const Tkey key = ptr->first; if(omp_test_lock(&lock_Ds_result_add_map.at(key))) { LRI_Cal_Aux::add_Ds(std::move(ptr->second), Ds_result.at(key), fac); @@ -183,11 +203,13 @@ namespace LRI_Cal_Aux } } - template + // Tkey labels thread lock, in actual use, Tkey can be TA, TC, Tk, etc. + // Tvalue can be Tensor or another map, both of them can be input of add_Ds + template void add_Ds_omp_wait_map( - std::map>> &Ds_result_thread, - std::map>> &Ds_result, - std::map &lock_Ds_result_add_map, + std::map& Ds_result_thread, + std::map& Ds_result, + std::map &lock_Ds_result_add_map, const double &fac) { if(Ds_result_thread.empty()) @@ -338,11 +360,11 @@ namespace LRI_Cal_Aux return list_filter; } - template + template std::map init_lock_result( const std::vector &labels, const std::unordered_map> &list_A, - std::map>> &Ds_result) + std::map>& Ds_result) { std::map lock_Ds_result_add_map; for(const Label::ab_ab &label : labels) @@ -375,10 +397,28 @@ namespace LRI_Cal_Aux return lock_Ds_result_add_map; } - template + template + std::map init_lock_result( + std::map& Ds_result, + std::vector key_list) + { + std::map lock_Ds_result_add_map; + + for(const Tkey &ikey : key_list) + { + Ds_result[ikey]; + lock_Ds_result_add_map[ikey]; + } + + for(auto &lock : lock_Ds_result_add_map) + omp_init_lock(&lock.second); + return lock_Ds_result_add_map; + } + + template void destroy_lock_result( - std::map &locks, - std::map &Ds_result) + std::map &locks, + std::map &Ds_result) { for(auto &lock : locks) omp_destroy_lock(&lock.second); diff --git a/include/RI/ri/LRI_Cal_Tools.h b/include/RI/ri/LRI_Cal_Tools.h index d9b0b0d..7d19954 100644 --- a/include/RI/ri/LRI_Cal_Tools.h +++ b/include/RI/ri/LRI_Cal_Tools.h @@ -58,8 +58,7 @@ class LRI_Cal_Tools *this->Ds_ab_ptr.at(label), Aa.first, TAC{Ab.first, (Ab.second-Aa.second)%this->period}); } - - std::vector split_b01(const Label::ab_ab &label) const + inline std::vector split_b01(const Label::ab_ab &label) const { switch(label) { diff --git a/include/RI/ri/LRI_k-cal_cvc_mo.hpp b/include/RI/ri/LRI_k-cal_cvc_mo.hpp new file mode 100644 index 0000000..f54975d --- /dev/null +++ b/include/RI/ri/LRI_k-cal_cvc_mo.hpp @@ -0,0 +1,502 @@ +// =================== +// Author: Ziqing Guan +// date: 2026.08.02 +// =================== + +#pragma once + +#include "LRI_k.h" +#include "LRI_Cal_Aux.h" +#include "../global/Array_Operator.h" +#include "../global/Tensor_Multiply.h" +#include +#include +#include +#ifdef __MKL_RI +#include +#endif + +namespace RI +{ +inline void switch_mo_type(const std::string &type, + std::size_t &imo, std::size_t &nmo, + std::size_t nocc, std::size_t nvirt) +{ + if (type == "O") + { + imo = 0; + nmo = nocc; + } + else if (type == "V") + { + imo = nocc; + nmo = nvirt; + } + else + { + throw std::runtime_error("Error in Cs_ao_mo_to_Cs_mo: unknown mo type " + type); + } +} + +/// @brief calculate Csk_ao_mo by C'^\mu (s,m)[k] = C^\mu (s,t)[k] c(m,t)[k] +/// s,t: atom orbital index; m: band index +template +std::map>> +LRI_k::cal_Csk_ao_mo( + const std::map>>& CsR_ao, // C^mu (s,t)[R] + const std::map>>& map_psi, // c(m,s)[k] + const std::vector& kindex_map, + const std::vector& k_indices, const std::vector& list_IJ, std::ofstream& ofs) +{ + using U64 = unsigned long long; + const std::size_t nmo = map_psi.begin()->second.begin()->second.shape[0]; + auto read_proc_status_kb = [](const std::string& key) -> U64 { + std::ifstream ifs("/proc/self/status"); + std::string line; + while (std::getline(ifs, line)) + { + if (line.compare(0, key.size(), key) == 0) + { + const std::size_t p = line.find(':'); + if (p == std::string::npos) return 0; + std::istringstream iss(line.substr(p + 1)); + U64 kb = 0; std::string unit; + iss >> kb >> unit; + return kb; + } + } + return 0; + }; + auto safe_mul = [](const U64 a, const U64 b) -> U64 { + if (a == 0 || b == 0) return 0; + if (a > std::numeric_limits::max() / b) + throw std::overflow_error("Multiplication would overflow"); + return a * b; + }; + U64 prealloc_count = 0; + U64 prealloc_est_bytes = 0; + U64 prealloc_max_bytes = 0; + int rank; + MPI_Comm_rank(this->mpi_comm, &rank); + // 1. allocate C^mu (s,m)[k] >> + std::map>> Csk_ao_mo; + for (const int ik : k_indices) + { + for (const auto& iat1 : list_IJ) + { + const std::size_t nabf = CsR_ao.at(iat1).begin()->second.shape[0]; + const std::size_t nw1 = CsR_ao.at(iat1).begin()->second.shape[1]; + const U64 tensor_bytes = safe_mul(safe_mul(static_cast(nabf), static_cast(nw1)), + safe_mul(static_cast(nmo), static_cast(sizeof(Tdata)))); + prealloc_count += 1; + prealloc_est_bytes += tensor_bytes; + if (tensor_bytes > prealloc_max_bytes) prealloc_max_bytes = tensor_bytes; + try + { + Csk_ao_mo[ik][iat1] = RI::Tensor({nabf, nw1, nmo}); // initialize + } + catch (const std::bad_alloc&) + { + const U64 vmrss_kb = read_proc_status_kb("VmRSS"); + const U64 vmhwm_kb = read_proc_status_kb("VmHWM"); + ofs << "[RI_MEMDBG] bad_alloc(prealloc): rank=" << rank + << ", tensors=" << prealloc_count + << ", ik=" << ik << ", iat1=" << iat1 + << ", shape={" << nabf << "," << nw1 << "," << nmo << "}" + << ", tensorMB=" << (tensor_bytes / 1024.0 / 1024.0) + << ", preallocGB=" << (prealloc_est_bytes / 1024.0 / 1024.0 / 1024.0) + << ", VmRSSGB=" << (vmrss_kb / 1024.0 / 1024.0) + << ", VmHWMGB=" << (vmhwm_kb / 1024.0 / 1024.0) + << std::endl; + throw; + } + } + } + + ofs << "Csk_ao_mo keys has been prepared." << std::endl; + const U64 vmrss_kb = read_proc_status_kb("VmRSS"); + const U64 vmhwm_kb = read_proc_status_kb("VmHWM"); + ofs << "[RI_MEMDBG] prealloc summary: rank=" << rank + << ", tensors=" << prealloc_count + << ", estGB=" << (prealloc_est_bytes / 1024.0 / 1024.0 / 1024.0) + << ", maxTensorMB=" << (prealloc_max_bytes / 1024.0 / 1024.0) + << ", VmRSSGB=" << (vmrss_kb / 1024.0 / 1024.0) + << ", VmHWMGB=" << (vmhwm_kb / 1024.0 / 1024.0) + << std::endl; + + // 2. calculate C'^mu (s,m)[k] = C^mu (s,t)[k] c(m,t)[k] +#ifdef __MKL_RI + const std::size_t mkl_threads = mkl_get_max_threads(); + mkl_set_num_threads(1); +#endif + +#pragma omp parallel for schedule(static) collapse(2) + for (const int ik : k_indices) + { + for (const int iat1 : list_IJ) + { + Tk k = kindex_map.at(ik); + auto& tensor_ao_mo = Csk_ao_mo.at(ik).at(iat1); // C'^\mu (s,m)[k] + auto& psi_k = map_psi.at(ik); + std::map> Ck_I_thread; + auto& CR_I = CsR_ao.at(iat1); + const std::size_t nabf = CR_I.begin()->second.shape[0]; + const std::size_t nw1 = CR_I.begin()->second.shape[1]; + // FT to get C^\mu (s,t)[k] = sum_R C^\mu (s,t)[R] exp(i k R) + for (const auto& CI_JR: CR_I) + { + const int iat2 = CI_JR.first.first; + const TC& R = CI_JR.first.second; + double arg = 2.0 * M_PI * (k[0] * R[0] + k[1] * R[1] + k[2] * R[2]); + std::complex phase(cos(arg), sin(arg)); + if (!Ck_I_thread.count(iat2)) + { + Ck_I_thread[iat2] = CI_JR.second * Global_Func::convert(phase); + } + else { LRI_Cal_Aux::add_Ds(CI_JR.second, Ck_I_thread[iat2], Global_Func::convert(phase)); } + } + // C'^mu (s,m)[k] = sum_t C^mu (s,t)[k] c(m,t)[k] + for (const auto& Ck_IJ: Ck_I_thread) + { + const int iat2 = Ck_IJ.first; + const auto& tensor_ao = Ck_IJ.second; // C^mu (s,t)[k] + const auto& psi_k_J = psi_k.at(iat2); // c(m,t)[k] + + const int nw2 = tensor_ao.shape[2]; + assert(nw2 == psi_k_J.shape[1]); + + // caution: Tensors are row-major (iw2 contiguous) + // C'(mu,s,m) = C(mu,s,t) c(m,t) << row-major + // C'_m_s_mu = (c_t_m)^T (C_t_s_mu) << col-major + // d/zgemm_('T', 'N', nmo, nw1*nabf, nw2, + // 1.0, psi_k_J.ptr(), nw2, + // tensor_ao.ptr(), nw2, + // 1.0, tensor_ao_mo.ptr(), nmo); + Blas_Interface::gemm( + 'N', 'T', + nw1*nabf, nmo, nw2, + Tdata(1.0), tensor_ao.ptr(), nw2, psi_k_J.ptr(), nw2, + Tdata(1.0), tensor_ao_mo.ptr(), nmo); + // tensor_ao_mo += Tensor_Multiply::x0x1y0_x0x1a_y0a(tensor_ao, psi_k_J); + } + } + } +#ifdef __MKL_RI + mkl_set_num_threads(mkl_threads); +#endif + return Csk_ao_mo; +} + +/// @brief left_conj: C^mu (m1^*,m2)[k1,k2] = c^*(m1,s)[k1] C^mu (s,m2)[k2] + C^*mu (s,m1)[k1] c(m2,s)[k2] +/// if right_conj: C^mu (m1,m2^*)[k1,k2] = c(m1,s)[k1] C^*mu (s,m2)[k2] + C^mu (s,m1)[k1] c^*(m2,s)[k2] +/// s: atom orbital index; m: band index +template +Tensor Cs_ao_mo_to_Cs_mo( + const std::map>>& Cs_ao_mo, // C^mu (s,m)[k] + const std::map>>& psi, // c(m,s)[k] + const int k1, const int k2, + TA mu, const std::string& m1_type, const std::string& m2_type, + const std::size_t nocc, const std::size_t nvirt, const bool is_left_conj) +{ + const Tensor& Cs_ao_mo_k1 = Cs_ao_mo.at(k1).at(mu); + const Tensor& Cs_ao_mo_k2 = Cs_ao_mo.at(k2).at(mu); + const Tensor& psi_k1 = psi.at(k1).at(mu); + const Tensor& psi_k2 = psi.at(k2).at(mu); + const std::size_t nabf = Cs_ao_mo_k1.shape[0]; + const std::size_t nw = Cs_ao_mo_k1.shape[1]; + const std::size_t nband = Cs_ao_mo_k1.shape[2]; + assert(nband == nocc + nvirt); + + std::size_t imo1, nmo1, imo2, nmo2; + switch_mo_type(m1_type, imo1, nmo1, nocc, nvirt); + switch_mo_type(m2_type, imo2, nmo2, nocc, nvirt); + Tensor Cs_mo({ nabf, nmo1, nmo2 }); + if (is_left_conj) + { + std::vector psi_k1_conj(nw * nmo1); + for (std::size_t m1 = 0; m1 < nmo1; ++m1) + for (std::size_t iw = 0; iw < nw; ++iw) + psi_k1_conj[m1 * nw + iw] = Global_Func::get_conj(psi_k1(imo1 + m1, iw)); + for (std::size_t iabf = 0; iabf < nabf; ++iabf) + { + Tdata *ptr_out = &Cs_mo(iabf, 0, 0); + int lda = nw; + int ldb = nband; + int ldc = nmo2; + Blas_Interface::gemm( + 'N', 'N', + nmo1, nmo2, nw, + Tdata(1.0), psi_k1_conj.data(), lda, &Cs_ao_mo_k2(iabf, 0, imo2), ldb, + Tdata(0.0), ptr_out, ldc); + lda = nband; + ldb = nw; + Blas_Interface::gemm( + 'C', 'T', + nmo1, nmo2, nw, + Tdata(1.0), &Cs_ao_mo_k1(iabf, 0, imo1), lda, &psi_k2(imo2, 0), ldb, + Tdata(1.0), ptr_out, ldc); + } + } + else + { + std::vector Cs_ao_mo_k2_conj(nabf * nw * nmo2); + for (std::size_t iabf = 0; iabf < nabf; ++iabf) + for (std::size_t iw = 0; iw < nw; ++iw) + for (std::size_t m2 = 0; m2 < nmo2; ++m2) + { + Cs_ao_mo_k2_conj[iabf * (nw * nmo2) + iw * nmo2 + m2] + = Global_Func::get_conj(Cs_ao_mo_k2(iabf, iw, imo2 + m2)); + } + + for (std::size_t iabf = 0; iabf < nabf; ++iabf) + { + Tdata *ptr_out = &Cs_mo(iabf, 0, 0); + int lda = nw; + int ldb = nmo2; + int ldc = nmo2; + Blas_Interface::gemm( + 'N', 'N', + nmo1, nmo2, nw, + Tdata(1.0), &psi_k1(imo1, 0), lda, &Cs_ao_mo_k2_conj[iabf * (nw * nmo2)], ldb, + Tdata(0.0), ptr_out, ldc); + lda = nband; + ldb = nw; + Blas_Interface::gemm( + 'T', 'C', + nmo1, nmo2, nw, + Tdata(1.0), &Cs_ao_mo_k1(iabf, 0, imo1), lda, &psi_k2(imo2, 0), ldb, + Tdata(1.0), ptr_out, ldc); + } + } + + return Cs_mo; +} + +template +std::map>> +LRI_k::cal_cvc_mo_k_onthefly( + const std::map>>& Cs_ao_mo, // C^mu (s,m)[k] + const std::map>>& map_psi, // c(m,s)[k] + const std::vector& k1_indices, + const std::vector& k2_indices, + const std::vector& list_I, + const std::vector& list_J, + const std::vector& psi_type, + const std::size_t nocc, + const std::size_t nvirt, + const std::string& save_name, + const bool is_A, + const std::vector& q_list, + const std::map>>& q2kpair) +{ + using namespace Array_Operator; +#ifdef __MKL_RI + const std::size_t mkl_threads = mkl_get_max_threads(); + mkl_set_num_threads(1); +#endif + std::map>> cvc_mo_k; + const std::map>>& Vs = this->data_pool.at(save_name).Ds_ab; + + std::map lock_cvc_result_add_map = LRI_Cal_Aux::init_lock_result(cvc_mo_k, k1_indices); +#pragma omp parallel + { + // 2. calculate CVC_mo_k + std::map>> cvc_mo_k_thread; +#pragma omp for schedule(dynamic, 64) collapse(3) + for (const Tk q: q_list) + { + for (const TA mu: list_I) + { + for (const TA nu: list_J) + { // 2.1 calculate V(q)_{mu,nu} on-the-fly + Tensor Vq_mu_nu; + const auto& V_mu = Vs.at(mu); + const std::vector list_nuR = Divide_Atoms::traversal_atom_period(std::vector{nu}, this->period); + for (const TAC& nu_R : list_nuR) + { + const Tensor& V_mu_nu_R = Global_Func::find(V_mu, nu_R); + if (V_mu_nu_R.empty()) continue; + const TC& R = nu_R.second; + double arg = 2.0 * M_PI * (q[0] * R[0] + q[1] * R[1] + q[2] * R[2]); + std::complex fac (cos(arg), sin(arg)); + LRI_Cal_Aux::add_Ds(V_mu_nu_R, Vq_mu_nu, Global_Func::convert(fac)); + } + if (Vq_mu_nu.empty()) continue; + + for (const auto& kpair: q2kpair.at(q)) + { + int k1 = kpair.first; + int k2 = kpair.second; + // 2.2 calculate C^mu (m1^*,m2)[k2,k1] on-the-fly, C_mu_ji for A and C_mu_bi for B + const Tensor C_mu_ji = Cs_ao_mo_to_Cs_mo(Cs_ao_mo, map_psi, k2, k1, + mu, psi_type[0], psi_type[1], nocc, nvirt, true); + // 2.3 calculate C^nu (m3,m4^*)[k2,k1] on-the-fly, C_nu_ba for A and C_nu_ja for B + const Tensor C_nu_ba = Cs_ao_mo_to_Cs_mo(Cs_ao_mo, map_psi, k2, k1, + nu, psi_type[2], psi_type[3], nocc, nvirt, false); + // 2.4 calculate CVC_mo + // CV_{ji,nu} = C^mu_{ji} V_{mu,nu} | CV_{bi,nu} = C^mu_{bi} V_{mu,nu} + const Tensor CV_ji_nu = Tensor_Multiply::x1x2y1_ax1x2_ay1(C_mu_ji, Vq_mu_nu); + const std::size_t nnu = Vq_mu_nu.shape[1]; + if (is_A) + { // (j,i nu) * (nu,b a) = (jiba) -> (j,bia) + // ̅ ̅ ̅ ̅ + Tensor cvc({nocc, nvirt, nocc, nvirt}); + for (std::size_t b = 0; b < nvirt; ++b) + for (std::size_t i = 0; i < nocc; ++i) + { + int lda = nnu * nocc; + int ldb = nvirt * nvirt; + int ldc = nvirt * nocc * nvirt; + Blas_Interface::gemm( + 'N', 'N', + nocc, nvirt, nnu, + Tdata(1.0), &CV_ji_nu(0, i, 0), lda, &C_nu_ba(0, b, 0), ldb, + Tdata(0.0), &cvc(0, b, i, 0), ldc); + } + LRI_Cal_Aux::add_Ds(std::move(cvc), cvc_mo_k_thread[k1][k2]); + } + else + { // (b i,nu) * (nu,j a) = (bija) -> (jb,ia) + // ̅ ̅ ̅ ̅ + Tensor cvc({nocc, nvirt, nocc, nvirt}); + for (std::size_t j = 0; j < nocc; ++j) + for (std::size_t i = 0; i < nocc; ++i) + { + int lda = nnu * nocc; + int ldb = nvirt * nocc; + int ldc = nvirt * nocc; + Blas_Interface::gemm( + 'N', 'N', + nvirt, nvirt, nnu,/*CV_bi_nu*/ /*C_nu_ja*/ + Tdata(1.0), &CV_ji_nu(0, i, 0), lda, &C_nu_ba(0, j, 0), ldb, + Tdata(0.0), &cvc(j, 0, i, 0), ldc); + } + LRI_Cal_Aux::add_Ds(std::move(cvc), cvc_mo_k_thread[k1][k2]); + } + } // end for kpair + LRI_Cal_Aux::add_Ds_omp_try_map(cvc_mo_k_thread, cvc_mo_k, lock_cvc_result_add_map, 1.0); + } // end for nu + } // end for mu + } // end for q + LRI_Cal_Aux::add_Ds_omp_wait_map(cvc_mo_k_thread, cvc_mo_k, lock_cvc_result_add_map, 1.0); + } // end #pragma omp parallel + + LRI_Cal_Aux::destroy_lock_result(lock_cvc_result_add_map, cvc_mo_k); + +#ifdef __MKL_RI + mkl_set_num_threads(mkl_threads); +#endif + + malloc_trim(0); + return cvc_mo_k; +} + +template +std::map>> +LRI_k::cal_cvc_mo_k_hartree_onthefly( + const std::map>>& Cs_ao_mo, // C^mu (s,m)[k] + const std::map>>& map_psi, // c(m,t)[k] + const std::vector& k1_indices, + const std::vector& k2_indices, + const std::vector& list_I, + const std::vector& list_J, + const std::vector& psi_type, + const std::size_t nocc, + const std::size_t nvirt, + const std::string& save_name, + const bool is_A) +{ +#ifdef __MKL_RI + const std::size_t mkl_threads = mkl_get_max_threads(); + mkl_set_num_threads(1); +#endif + + std::map>> cvc_mo_k; + std::map>> Vq; // has only one q=0, keep for further q != 0 extension + + const std::map>>& Vs = this->data_pool.at(save_name).Ds_ab; + + // add thread lock for the first Tk key of cvc_mo_k + std::map lock_cvc_result_add_map = LRI_Cal_Aux::init_lock_result(cvc_mo_k, k1_indices); + // add thread lock for the TA key of Vq + std::map lock_vq_result_add_map = LRI_Cal_Aux::init_lock_result(Vq, list_I); + const std::vector list_JR = Divide_Atoms::traversal_atom_period(list_J, this->period); + #pragma omp parallel + { + Tk q{0.0, 0.0, 0.0}; + // 1. FT V_mu_nu > to V_mu_nu > + std::map>> Vq_thread; + for (const TA mu : list_I) + { + const auto& V_mu = Vs.at(mu); + auto& Vq_mu_thread = Vq_thread[mu]; + #pragma omp for schedule(dynamic) + for (const TAC& nu_R : list_JR) + { + const Tensor& V_mu_nu_R = Global_Func::find(V_mu, nu_R); + if (V_mu_nu_R.empty()) continue; + const TA nu = nu_R.first; + const TC& R = nu_R.second; + double arg = 2.0 * M_PI * (q[0] * R[0] + q[1] * R[1] + q[2] * R[2]); + std::complex fac (cos(arg), sin(arg)); + LRI_Cal_Aux::add_Ds(V_mu_nu_R, Vq_mu_thread[nu], Global_Func::convert(fac)); + } + LRI_Cal_Aux::add_Ds_omp_try_map(Vq_thread, Vq, lock_vq_result_add_map, 1.0); + } + LRI_Cal_Aux::add_Ds_omp_wait_map(Vq_thread, Vq, lock_vq_result_add_map, 1.0); + + #pragma omp barrier + #pragma omp master + { + LRI_Cal_Aux::destroy_lock_result(lock_vq_result_add_map, Vq); + } + #pragma omp barrier + + // 2 calculate CVC_mo_k + std::map>> cvc_mo_k_thread; +#pragma omp for schedule(dynamic, 64) collapse(2) + for (const TA mu : list_I) + { + for (const int k1: k1_indices) + { + // 2.1 calculate C^mu (i,a^*)[k1,k1] on-the-fly + const Tensor C_mu_ia = Cs_ao_mo_to_Cs_mo(Cs_ao_mo, map_psi, k1, k1, + mu, psi_type[0], psi_type[1], nocc, nvirt, false); + for (const TA nu : list_J) + { + const Tensor& Vq_mu_nu = Global_Func::find(Vq, mu, nu); + if (Vq_mu_nu.empty()) continue; + for (const int k2: k2_indices) + { + // 2.2 calculate C^nu (m3,m4)[k2,k2] on-the-fly, C_nu_j^*b for A and C_nu_jb^* for B + Tensor C_nu_jb; + if (is_A){ + C_nu_jb = Cs_ao_mo_to_Cs_mo(Cs_ao_mo, map_psi, k2, k2, + nu, psi_type[2], psi_type[3], nocc, nvirt, true); + } + else{ + C_nu_jb = Cs_ao_mo_to_Cs_mo(Cs_ao_mo, map_psi, k2, k2, + nu, psi_type[2], psi_type[3], nocc, nvirt, false); + } + // 2.3 calculate CVC_mo + const Tensor CV_ia_nu = Tensor_Multiply::x1x2y1_ax1x2_ay1(C_mu_ia, Vq_mu_nu); + LRI_Cal_Aux::add_Ds(Tensor_Multiply::x1x2y0y1_ax1x2_y0y1a(C_nu_jb, CV_ia_nu), cvc_mo_k_thread[k1][k2]); + } + } + LRI_Cal_Aux::add_Ds_omp_try_map(cvc_mo_k_thread, cvc_mo_k, lock_cvc_result_add_map, 1.0); + } // end for k1 + } // end for mu + LRI_Cal_Aux::add_Ds_omp_wait_map(cvc_mo_k_thread, cvc_mo_k, lock_cvc_result_add_map, 1.0); + } // end #pragma omp parallel + + LRI_Cal_Aux::destroy_lock_result(lock_cvc_result_add_map, cvc_mo_k); + + #ifdef __MKL_RI + mkl_set_num_threads(mkl_threads); + #endif + + malloc_trim(0); + return cvc_mo_k; +} + +} // end namespace RI diff --git a/include/RI/ri/LRI_k-cal_hartree.hpp b/include/RI/ri/LRI_k-cal_hartree.hpp new file mode 100644 index 0000000..9863a0a --- /dev/null +++ b/include/RI/ri/LRI_k-cal_hartree.hpp @@ -0,0 +1,238 @@ +// =================== +// Author: Ziqing Guan +// date: 2026.08.02 +// =================== + +#pragma once + +#include "LRI_k.h" +#include "LRI_Cal_Aux.h" +#include "../global/Array_Operator.h" +#include "../global/Tensor_Multiply.h" +#include +#include +#include +#ifdef __MKL_RI +#include +#endif + +namespace RI +{ +template +std::map>>> // H(s,t)[k] +LRI_k::cal_cvcd_k_hartree( + const std::map>>>& Ds, // D(s,t)[k] + const std::vector& kindex_map,// k index to direct coordinate array + const std::vector& k_indices, + const std::vector& list_I, + const std::vector& list_J, + const std::vector& list_IJ, + const std::string& save_name_C, + const std::string& save_name_V) +{ +#ifdef __MKL_RI + const std::size_t mkl_threads = mkl_get_max_threads(); + mkl_set_num_threads(1); +#endif + const int nk = static_cast(kindex_map.size()); + assert (nk == this->period[0] * this->period[1] * this->period[2]); + std::map>>> hartree_k; + std::map>>> Csk; + std::map>> Vq; // has only one q=0 + std::map> M_nu; // M_nu = \sum_{uvk} (C^nu_u_v[k] + C^nu*_v_u[k]) D_v_u[k] + + const std::map>>& Vs = this->data_pool.at(save_name_V).Ds_ab; + const std::map>>& Cs = this->data_pool.at(save_name_C).Ds_ab; + + // add thread lock for the TA key of Vq + std::map lock_vq_result_add_map = LRI_Cal_Aux::init_lock_result(Vq, list_IJ); + const std::vector list_IJR = Divide_Atoms::traversal_atom_period(list_IJ, this->period); + // add thread lock for the TA key of Csk + std::map lock_csk_result_add_map = LRI_Cal_Aux::init_lock_result(Csk, list_IJ); + + // add thread lock for the TA key of M_nu + std::map lock_m_result_add_map = LRI_Cal_Aux::init_lock_result(M_nu, list_IJ); + #pragma omp parallel + { + Tk q{0.0, 0.0, 0.0}; + // 1. FT V_mu_nu > to V_mu_nu > + std::map>> Vq_thread; + + for (const TA mu : list_I) + { + const auto& V_mu = Vs.at(mu); + auto& Mu_Vq_nu_thread = Vq_thread[mu]; + #pragma omp for schedule(dynamic) nowait + for (const TAC& nu_R : list_IJR) + { + const Tensor& V_mu_nu_R = Global_Func::find(V_mu, nu_R); + if (V_mu_nu_R.empty()) continue; + const TA nu = nu_R.first; + const TC R = nu_R.second; + double arg = 2.0 * M_PI * (q[0] * R[0] + q[1] * R[1] + q[2] * R[2]); + std::complex fac (cos(arg), sin(arg)); + LRI_Cal_Aux::add_Ds(V_mu_nu_R, Mu_Vq_nu_thread[nu], Global_Func::convert(fac)); + } + LRI_Cal_Aux::add_Ds_omp_try_map(Vq_thread, Vq, lock_vq_result_add_map, 1.0); + } + + // 2. FT CsR > to Csk >> + std::map>>> Csk_thread; + for (const TA mu : list_IJ) + { + const auto& Cs_mu = Cs.at(mu); + auto& Mu_C_nu_k_thread = Csk_thread[mu]; + #pragma omp for schedule(dynamic) nowait + for (const TAC& nu_R : list_IJR) + { + const Tensor& Cs_mu_nu_R = Global_Func::find(Cs_mu, nu_R); + if (Cs_mu_nu_R.empty()) continue; + const TA nu = nu_R.first; + const TC R = nu_R.second; + auto& Nu_C_k_thread = Mu_C_nu_k_thread[nu]; + for (const int ik: k_indices) + { + const Tk k = kindex_map.at(ik); + double arg = 2.0 * M_PI * (k[0] * R[0] + k[1] * R[1] + k[2] * R[2]); + std::complex fac (cos(arg), sin(arg)); + LRI_Cal_Aux::add_Ds(Cs_mu_nu_R, Nu_C_k_thread[ik], Global_Func::convert(fac)); + } + } + LRI_Cal_Aux::add_Ds_omp_try_map(Csk_thread, Csk, lock_csk_result_add_map, 1.0); + } + LRI_Cal_Aux::add_Ds_omp_wait_map(Vq_thread, Vq, lock_vq_result_add_map, 1.0); + LRI_Cal_Aux::add_Ds_omp_wait_map(Csk_thread, Csk, lock_csk_result_add_map, 1.0); + #pragma omp barrier + #pragma omp master + { + LRI_Cal_Aux::destroy_lock_result(lock_vq_result_add_map, Vq); + this->free_tensors_map2(save_name_V); + LRI_Cal_Aux::destroy_lock_result(lock_csk_result_add_map, Csk); + this->free_tensors_map2(save_name_C); + } + #pragma omp barrier + + // 3. calculate M_nu =\sum_{uvk} (C^nu_u_v[k] + C^nu*_v_u[k]) D_v_u[k] + std::map> M_nu_thread; +#pragma omp for schedule(dynamic) collapse(2) + for (const TA v : list_I) + { + for (const TA u : list_J) + { + const std::map>& D_v_u = Global_Func::find(Ds, v, u); + if (D_v_u.empty()) continue; + const std::size_t nwt1 = D_v_u.begin()->second.shape[0]; + const std::size_t nwt2 = D_v_u.begin()->second.shape[1]; + const std::map>& Csk_u_v = Global_Func::find(Csk, u, v); + const std::map>& Csk_v_u = Global_Func::find(Csk, v, u); + for (const int ik: k_indices) + { + if (!Csk_u_v.empty()) + { + const Tensor& D_v_u_k = Global_Func::find(D_v_u, ik); + if (D_v_u_k.empty()) continue; + const Tensor& C_u_v_k = Global_Func::find(Csk_u_v, ik); + if (C_u_v_k.empty()) continue; + assert(C_u_v_k.shape[1]==nwt2); + assert(C_u_v_k.shape[2]==nwt1); + // M_nu += C^nu_u_v[k] * D_v_u[k] + LRI_Cal_Aux::add_Ds(Tensor_Multiply::gemv(C_u_v_k, D_v_u_k.transpose()), M_nu_thread[u]); + } + if (!Csk_v_u.empty()) + { + const Tensor& D_v_u_k = Global_Func::find(D_v_u, ik); + if (D_v_u_k.empty()) continue; + const Tensor& C_v_u_k = Global_Func::find(Csk_v_u, ik); + if (C_v_u_k.empty()) continue; + assert(C_v_u_k.shape[1]==nwt1); + assert(C_v_u_k.shape[2]==nwt2); + // M_nu += C^nu*_v_u[k] * D_v_u[k] + LRI_Cal_Aux::add_Ds(Tensor_Multiply::gemv(C_v_u_k.conjugate(), D_v_u_k), M_nu_thread[v]); + } + } + } + } + LRI_Cal_Aux::add_Ds_omp_wait_map(M_nu_thread, M_nu, lock_m_result_add_map, 1.0); + }// end #pragma omp parallel + LRI_Cal_Aux::destroy_lock_result(lock_m_result_add_map, M_nu); +#ifdef __MKL_RI + mkl_set_num_threads(mkl_threads); +#endif + + // 4. calculate N_mu = 1/Nk \sum_nu Vq_mu_nu * M_nu, and communicate N_mu + std::map> N_mu; + for (int mu : list_I) + { + const auto& Vq_mu = Vq.at(mu); + Tensor& N_Mu = N_mu[mu]; + for (const TA nu : list_IJ) + { + const Tensor& Vq_mu_nu = Global_Func::find(Vq_mu, nu); + if (Vq_mu_nu.empty()) continue; + const Tensor& M_Nu = M_nu.at(nu); + LRI_Cal_Aux::add_Ds(Tensor_Multiply::gemv(Vq_mu_nu, M_Nu), N_Mu); + } + const Tdata fac = Tdata(1.0 / static_cast(nk)); + N_Mu = fac * N_Mu; + } + std::set set_IJ(list_IJ.begin(), list_IJ.end()); + N_mu = Communicate_Tensors_Map_Judge::comm_map(this->mpi_comm, std::move(N_mu), set_IJ); + + // 5. calculate H_st[k] = \sum_mu (C^mu_s_t[k] + C^mu*_t_s[k]) N_mu + for (const TA s : list_I) + { + std::map>>& hartree_s = hartree_k[s]; + for (const TA t : list_J) + { + std::map>& hartree_st = hartree_s[t]; + for (const int ik: k_indices) + { + hartree_st[ik]; + } + } + } +#ifdef __MKL_RI + mkl_set_num_threads(1); +#endif + +#pragma omp parallel for schedule(dynamic) collapse(2) + for (const TA s : list_I) + { + for (const TA t : list_J) + { + auto& hartree_st = hartree_k.at(s).at(t); + const std::map>& Csk_s_t = Global_Func::find(Csk, s, t); + const std::map>& Csk_t_s = Global_Func::find(Csk, t, s); + const Tensor& N_Mu_at_s = N_mu.at(s); + const Tensor& N_Mu_at_t = N_mu.at(t); + + for (const int ik: k_indices) + { + if (!Csk_s_t.empty()) + { + const Tensor& C_s_t_k = Global_Func::find(Csk_s_t, ik); + if (C_s_t_k.empty()) continue; + // H_st[k] += C^mu_s_t[k] * N_mu + LRI_Cal_Aux::add_Ds(Tensor_Multiply::gemv_trans(C_s_t_k, N_Mu_at_s), hartree_st.at(ik)); + } + if (!Csk_t_s.empty()) + { + const Tensor& C_t_s_k = Global_Func::find(Csk_t_s, ik); + if (C_t_s_k.empty()) continue; + // H_st[k] += C^mu*_t_s[k] * N_mu + Tensor tmp = Tensor_Multiply::gemv_trans(C_t_s_k.conjugate(), N_Mu_at_t); + LRI_Cal_Aux::add_Ds(tmp.transpose(), hartree_st.at(ik)); + } + } + } + } + +#ifdef __MKL_RI + mkl_set_num_threads(mkl_threads); +#endif + + malloc_trim(0); + return hartree_k; +} + +} // end namespace RI diff --git a/include/RI/ri/LRI_k.h b/include/RI/ri/LRI_k.h new file mode 100644 index 0000000..2450283 --- /dev/null +++ b/include/RI/ri/LRI_k.h @@ -0,0 +1,72 @@ +// =================== +// Author: Ziqing Guan +// date: 2026.08.02 +// =================== + +#pragma once + +#include "LRI.h" + +namespace RI +{ + +template +class LRI_k : public LRI +{ +public: + using TC = std::array; + using TAC = std::pair; + using Tdata_real = Global_Func::To_Real_t; + using Tk = std::array; + + LRI_k() = default; + + std::map>> cal_Csk_ao_mo( + const std::map>>& CsR_ao, + const std::map>>& map_psi, + const std::vector& kindex_map, + const std::vector& k_indices, const std::vector& list_IJ, std::ofstream& ofs); + + std::map>> cal_cvc_mo_k_onthefly( + const std::map>>& Cs_ao_mo, + const std::map>>& map_psi, + const std::vector& k1_indices, + const std::vector& k2_indices, + const std::vector& list_I, + const std::vector& list_J, + const std::vector& psi_type, + const std::size_t nocc, + const std::size_t nvirt, + const std::string& save_name, + const bool is_A, + const std::vector& q_list_in, + const std::map>>& q2kpair_in); + + std::map>> cal_cvc_mo_k_hartree_onthefly( + const std::map>>& Cs_ao_mo, + const std::map>>& map_psi, + const std::vector& k1_indices, + const std::vector& k2_indices, + const std::vector& list_I, + const std::vector& list_J, + const std::vector& psi_type, + const std::size_t nocc, + const std::size_t nvirt, + const std::string& save_name, + const bool is_A); + + std::map>>> cal_cvcd_k_hartree( + const std::map>>>& Ds, // D(s,t)[k] + const std::vector& kindex_map,// k index to direct coordinate array + const std::vector& list_k_index, + const std::vector& list_I, + const std::vector& list_J, + const std::vector& list_IJ, + const std::string& save_name_C, + const std::string& save_name_V); + +}; +} + +#include "LRI_k-cal_cvc_mo.hpp" +#include "LRI_k-cal_hartree.hpp" diff --git a/unittests/Test_All.hpp b/unittests/Test_All.hpp index e34f3e1..70fd32a 100644 --- a/unittests/Test_All.hpp +++ b/unittests/Test_All.hpp @@ -21,6 +21,8 @@ #include "unittests/physics/Exx-test.hpp" #include "unittests/physics/RPA-test.hpp" #include "unittests/physics/GW-test.hpp" +#include "unittests/physics/LR-test.hpp" +#include "unittests/physics/Hartree-test.hpp" namespace Test_All { @@ -92,5 +94,8 @@ namespace Test_All GW_Test::main(argc, argv); GW_Test::main>(argc, argv); GW_Test::main>(argc, argv); + + LR_Test::main(argc, argv); + Hartree_Test::main(argc, argv); } } \ No newline at end of file diff --git a/unittests/physics/Hartree-ref b/unittests/physics/Hartree-ref new file mode 100644 index 0000000..4dbda44 --- /dev/null +++ b/unittests/physics/Hartree-ref @@ -0,0 +1,6 @@ +1 +0 1 +0 3 +0 2 3 3 3.292800000e+01 1.355200000e+01 6.585600000e+01 2.710400000e+01 9.878400000e+01 4.065600000e+01 6.585600000e+01 2.710400000e+01 1.317120000e+02 5.420800000e+01 1.975680000e+02 8.131200000e+01 9.878400000e+01 4.065600000e+01 1.975680000e+02 8.131200000e+01 2.963520000e+02 1.219680000e+02 +1 2 3 3 8.232000000e+00 3.388000000e+00 1.646400000e+01 6.776000000e+00 2.469600000e+01 1.016400000e+01 1.646400000e+01 6.776000000e+00 3.292800000e+01 1.355200000e+01 4.939200000e+01 2.032800000e+01 2.469600000e+01 1.016400000e+01 4.939200000e+01 2.032800000e+01 7.408800000e+01 3.049200000e+01 +2 2 3 3 8.232000000e+00 3.388000000e+00 1.646400000e+01 6.776000000e+00 2.469600000e+01 1.016400000e+01 1.646400000e+01 6.776000000e+00 3.292800000e+01 1.355200000e+01 4.939200000e+01 2.032800000e+01 2.469600000e+01 1.016400000e+01 4.939200000e+01 2.032800000e+01 7.408800000e+01 3.049200000e+01 diff --git a/unittests/physics/Hartree-test.hpp b/unittests/physics/Hartree-test.hpp new file mode 100644 index 0000000..1b4a0f0 --- /dev/null +++ b/unittests/physics/Hartree-test.hpp @@ -0,0 +1,109 @@ +// =================== +// Author: Ziqing Guan +// date: 2026.08.02 +// =================== + +#pragma once + +#include "RI/physics/Hartree.h" +#include "Test_Helpers.hpp" + +namespace Hartree_Test +{ + +using namespace Test_Helpers; +static const std::string ref_path = "unittests/physics/Hartree-ref"; + +// -------- main test -------- +static void main(int argc, char *argv[]) +{ + int mpi_init_provide; + MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &mpi_init_provide); + + // ====== config: 1 atom, 3x1x1 k-grid, nabf=2, nw=3 ====== + const TA atom0 = 0; + const int nk = 3; + const int nabf = 2, nw = 3; + + // 3x1x1 uniform k-grid fractional coordinates + std::vector kfrac(nk); + for (int i = 0; i < nk; ++i) + kfrac[i] = {{double(i)/3.0, 0.0, 0.0}}; + + // Cs: shape {nabf=2, nw=3, nw=3}, values vary with R + std::map>> Cs; + for (auto R : {std::array{0,0,0}, std::array{1,0,0}, std::array{-1,0,0}}) { + const double w = (R[0]==0 && R[1]==0 && R[2]==0) ? 1.0 : 0.5 / (std::abs(R[0])+std::abs(R[1])+std::abs(R[2])); + Cs[atom0][{atom0, R}] = RI::Tensor({nabf, nw, nw}); + for (int mu = 0; mu < nabf; ++mu) + for (int s = 0; s < nw; ++s) + for (int t = 0; t < nw; ++t) + Cs[atom0][{atom0, R}](mu,s,t) = Tdata(w * (mu+1) * (s+1) * (t+1), 0.0); + } + + // Vs: shape {nabf=2, nabf=2}, values vary with R + std::map>> Vs; + for (auto R : {std::array{0,0,0}, std::array{1,0,0}, std::array{-1,0,0}}) { + const double w = (R[0]==0 && R[1]==0 && R[2]==0) ? 0.01 : 0.005 / (std::abs(R[0])+std::abs(R[1])+std::abs(R[2])); + Vs[atom0][{atom0, R}] = RI::Tensor({nabf, nabf}); + for (int i = 0; i < nabf; ++i) + for (int j = 0; j < nabf; ++j) + Vs[atom0][{atom0, R}](i,j) = (i==j) ? Tdata(w, 0.0) : Tdata(w*0.5, 0.0); + } + + // Ds: D(s,t)[k] for all k, shape {nw=3, nw=3}, complex values + std::map>>> Ds; + for (int ik = 0; ik < nk; ++ik) { + Ds[atom0][atom0][ik] = RI::Tensor({nw, nw}); + for (int i = 0; i < nw; ++i) + for (int j = 0; j < nw; ++j) + Ds[atom0][atom0][ik](i,j) + = Tdata(0.1 * (i+1.0)*(j+1.0)*(ik+1.0), + 0.1 * (i+0.5)*(j+0.5)*(ik+0.5)); + } + + // ====== init Hartree with MPI distribution ====== + RI::Hartree hartree; + hartree.init(kfrac); + hartree.set_parallel(MPI_COMM_WORLD, 1, nk, {{3,1,1}}); + hartree.set_Cs(Cs, 0.0, {atom0}, {atom0}); + hartree.set_Vs(Vs, 0.0, {atom0}, {atom0}); + + // compute + auto result = hartree.cal_hartree(Ds); + + // ====== gather per-rank results to rank 0, then regression ====== + int nproc; MPI_Comm_size(MPI_COMM_WORLD, &nproc); + auto gathered = gather_map(result, 0, MPI_COMM_WORLD, 0); + + if (nproc > 1) MPI_Barrier(MPI_COMM_WORLD); + + if (argc >= 2 && std::string(argv[1]) == "--write-ref") + { + assert(nproc == 1 && "--write-ref requires exactly 1 MPI process"); + std::ofstream ofs_ref(ref_path); + assert(ofs_ref && "Failed to open ref file for writing"); + ofs_ref << std::scientific << std::setprecision(REF_PRECISION); + write_map(ofs_ref, gathered); + std::cout << "Hartree_ref written to " << ref_path << std::endl; + } + else + { + int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank); + if (rank == 0) + { + std::ifstream ifs_ref(ref_path); + assert(ifs_ref && "Hartree_ref not found; run with --write-ref to generate"); + HartreeMap ref; + read_map(ifs_ref, ref); + compare_map(gathered, ref); + } + } + + hartree.free_Cs(); + hartree.free_Vs(); + + MPI_Finalize(); +} + +} diff --git a/unittests/physics/LR-ref b/unittests/physics/LR-ref new file mode 100644 index 0000000..b879027 --- /dev/null +++ b/unittests/physics/LR-ref @@ -0,0 +1,59 @@ +3 +0 1 +0 3 2 3 3 2.800000000e+00 5.500000000e-01 5.600000000e+00 1.650000000e+00 8.400000000e+00 2.750000000e+00 5.600000000e+00 1.100000000e+00 1.120000000e+01 3.300000000e+00 1.680000000e+01 5.500000000e+00 8.400000000e+00 1.650000000e+00 1.680000000e+01 4.950000000e+00 2.520000000e+01 8.250000000e+00 5.600000000e+00 1.100000000e+00 1.120000000e+01 3.300000000e+00 1.680000000e+01 5.500000000e+00 1.120000000e+01 2.200000000e+00 2.240000000e+01 6.600000000e+00 3.360000000e+01 1.100000000e+01 1.680000000e+01 3.300000000e+00 3.360000000e+01 9.900000000e+00 5.040000000e+01 1.650000000e+01 +1 1 +0 3 2 3 3 5.600000000e+00 1.650000000e+00 1.120000000e+01 4.950000000e+00 1.680000000e+01 8.250000000e+00 1.120000000e+01 3.300000000e+00 2.240000000e+01 9.900000000e+00 3.360000000e+01 1.650000000e+01 1.680000000e+01 4.950000000e+00 3.360000000e+01 1.485000000e+01 5.040000000e+01 2.475000000e+01 1.120000000e+01 3.300000000e+00 2.240000000e+01 9.900000000e+00 3.360000000e+01 1.650000000e+01 2.240000000e+01 6.600000000e+00 4.480000000e+01 1.980000000e+01 6.720000000e+01 3.300000000e+01 3.360000000e+01 9.900000000e+00 6.720000000e+01 2.970000000e+01 1.008000000e+02 4.950000000e+01 +2 1 +0 3 2 3 3 8.400000000e+00 2.750000000e+00 1.680000000e+01 8.250000000e+00 2.520000000e+01 1.375000000e+01 1.680000000e+01 5.500000000e+00 3.360000000e+01 1.650000000e+01 5.040000000e+01 2.750000000e+01 2.520000000e+01 8.250000000e+00 5.040000000e+01 2.475000000e+01 7.560000000e+01 4.125000000e+01 1.680000000e+01 5.500000000e+00 3.360000000e+01 1.650000000e+01 5.040000000e+01 2.750000000e+01 3.360000000e+01 1.100000000e+01 6.720000000e+01 3.300000000e+01 1.008000000e+02 5.500000000e+01 5.040000000e+01 1.650000000e+01 1.008000000e+02 4.950000000e+01 1.512000000e+02 8.250000000e+01 +3 +0 3 +0 4 1 2 1 2 3.885234588e+01 0.000000000e+00 5.879577113e+01 -1.755523000e+00 5.879577113e+01 1.755523000e+00 8.905574388e+01 2.220446049e-15 +1 4 1 2 1 2 1.666106409e+02 6.174245000e+00 2.524130761e+02 1.815352000e+00 2.540731961e+02 1.969667700e+01 3.853823269e+02 1.832707800e+01 +2 4 1 2 1 2 3.853823269e+02 1.832707800e+01 5.840322941e+02 1.032131100e+01 5.893446781e+02 5.349736700e+01 8.942803809e+02 5.432904400e+01 +1 3 +0 4 1 2 1 2 1.666106409e+02 -6.174245000e+00 2.540731961e+02 -1.969667700e+01 2.524130761e+02 -1.815352000e+00 3.853823269e+02 -1.832707800e+01 +1 4 1 2 1 2 7.154581359e+02 -1.421085472e-14 1.092673021e+03 -4.408912200e+01 1.092673021e+03 4.408912200e+01 1.671485894e+03 4.263256415e-14 +2 4 1 2 1 2 1.655548742e+03 1.734879300e+01 2.529481807e+03 -7.552521900e+01 2.535790263e+03 1.357568520e+02 3.881114588e+03 5.106809400e+01 +2 3 +0 4 1 2 1 2 3.853823269e+02 -1.832707800e+01 5.893446781e+02 -5.349736700e+01 5.840322941e+02 -1.032131100e+01 8.942803809e+02 -5.432904400e+01 +1 4 1 2 1 2 1.655548742e+03 -1.734879300e+01 2.535790263e+03 -1.357568520e+02 2.529481807e+03 7.552521900e+01 3.881114588e+03 -5.106809400e+01 +2 4 1 2 1 2 3.831310988e+03 -1.136868377e-13 5.871034777e+03 -2.526481650e+02 5.871034777e+03 2.526481650e+02 9.013332658e+03 0.000000000e+00 +3 +0 3 +0 4 1 2 1 2 3.818829788e+01 -7.152530000e+00 5.746767513e+01 -1.254953700e+01 5.746767513e+01 -1.254953700e+01 8.639955188e+01 -2.158802800e+01 +1 4 1 2 1 2 1.626263529e+02 -3.674093500e+01 2.444445001e+02 -6.294873200e+01 2.461046201e+02 -6.613368300e+01 3.694451749e+02 -1.112010900e+02 +2 4 1 2 1 2 3.754216069e+02 -8.896087200e+01 5.641108541e+02 -1.515888990e+02 5.694232381e+02 -1.610785330e+02 8.544375009e+02 -2.694913760e+02 +1 3 +0 4 1 2 1 2 1.626263529e+02 -3.674093500e+01 2.461046201e+02 -6.613368300e+01 2.444445001e+02 -6.294873200e+01 3.694451749e+02 -1.112010900e+02 +1 4 1 2 1 2 6.915524079e+02 -1.834001400e+02 1.044861565e+03 -3.227111580e+02 1.044861565e+03 -3.227111580e+02 1.575862982e+03 -5.572440720e+02 +2 4 1 2 1 2 1.595784422e+03 -4.411515570e+02 2.409953167e+03 -7.720803090e+02 2.416261623e+03 -7.812438480e+02 3.642057308e+03 -1.342042086e+03 +2 3 +0 4 1 2 1 2 3.754216069e+02 -8.896087200e+01 5.694232381e+02 -1.610785330e+02 5.641108541e+02 -1.515888990e+02 8.544375009e+02 -2.694913760e+02 +1 4 1 2 1 2 1.595784422e+03 -4.411515570e+02 2.416261623e+03 -7.812438480e+02 2.409953167e+03 -7.720803090e+02 3.642057308e+03 -1.342042086e+03 +2 4 1 2 1 2 3.681900188e+03 -1.059506910e+03 5.572213177e+03 -1.866365655e+03 5.572213177e+03 -1.866365655e+03 8.415689458e+03 -3.227434980e+03 +3 +0 3 +0 4 1 2 1 2 3.885234588e+01 0.000000000e+00 5.879577113e+01 -1.755523000e+00 5.879577113e+01 1.755523000e+00 8.905574388e+01 0.000000000e+00 +1 4 1 2 1 2 1.666106409e+02 6.174245000e+00 2.524130761e+02 1.815352000e+00 2.540731961e+02 1.969667700e+01 3.853823269e+02 1.832707800e+01 +2 4 1 2 1 2 3.853823269e+02 1.832707800e+01 5.840322941e+02 1.032131100e+01 5.893446781e+02 5.349736700e+01 8.942803809e+02 5.432904400e+01 +1 3 +0 4 1 2 1 2 1.666106409e+02 -6.174245000e+00 2.540731961e+02 -1.969667700e+01 2.524130761e+02 -1.815352000e+00 3.853823269e+02 -1.832707800e+01 +1 4 1 2 1 2 7.154581359e+02 0.000000000e+00 1.092673021e+03 -4.408912200e+01 1.092673021e+03 4.408912200e+01 1.671485894e+03 0.000000000e+00 +2 4 1 2 1 2 1.655548742e+03 1.734879300e+01 2.529481807e+03 -7.552521900e+01 2.535790263e+03 1.357568520e+02 3.881114588e+03 5.106809400e+01 +2 3 +0 4 1 2 1 2 3.853823269e+02 -1.832707800e+01 5.893446781e+02 -5.349736700e+01 5.840322941e+02 -1.032131100e+01 8.942803809e+02 -5.432904400e+01 +1 4 1 2 1 2 1.655548742e+03 -1.734879300e+01 2.535790263e+03 -1.357568520e+02 2.529481807e+03 7.552521900e+01 3.881114588e+03 -5.106809400e+01 +2 4 1 2 1 2 3.831310988e+03 0.000000000e+00 5.871034777e+03 -2.526481650e+02 5.871034777e+03 2.526481650e+02 9.013332658e+03 0.000000000e+00 +3 +0 3 +0 4 1 2 1 2 3.818829788e+01 -7.152530000e+00 5.746767513e+01 -1.254953700e+01 5.746767513e+01 -1.254953700e+01 8.639955188e+01 -2.158802800e+01 +1 4 1 2 1 2 1.626263529e+02 -3.674093500e+01 2.444445001e+02 -6.294873200e+01 2.461046201e+02 -6.613368300e+01 3.694451749e+02 -1.112010900e+02 +2 4 1 2 1 2 3.754216069e+02 -8.896087200e+01 5.641108541e+02 -1.515888990e+02 5.694232381e+02 -1.610785330e+02 8.544375009e+02 -2.694913760e+02 +1 3 +0 4 1 2 1 2 1.626263529e+02 -3.674093500e+01 2.461046201e+02 -6.613368300e+01 2.444445001e+02 -6.294873200e+01 3.694451749e+02 -1.112010900e+02 +1 4 1 2 1 2 6.915524079e+02 -1.834001400e+02 1.044861565e+03 -3.227111580e+02 1.044861565e+03 -3.227111580e+02 1.575862982e+03 -5.572440720e+02 +2 4 1 2 1 2 1.595784422e+03 -4.411515570e+02 2.409953167e+03 -7.720803090e+02 2.416261623e+03 -7.812438480e+02 3.642057308e+03 -1.342042086e+03 +2 3 +0 4 1 2 1 2 3.754216069e+02 -8.896087200e+01 5.694232381e+02 -1.610785330e+02 5.641108541e+02 -1.515888990e+02 8.544375009e+02 -2.694913760e+02 +1 4 1 2 1 2 1.595784422e+03 -4.411515570e+02 2.416261623e+03 -7.812438480e+02 2.409953167e+03 -7.720803090e+02 3.642057308e+03 -1.342042086e+03 +2 4 1 2 1 2 3.681900188e+03 -1.059506910e+03 5.572213177e+03 -1.866365655e+03 5.572213177e+03 -1.866365655e+03 8.415689458e+03 -3.227434980e+03 diff --git a/unittests/physics/LR-test.hpp b/unittests/physics/LR-test.hpp new file mode 100644 index 0000000..76b51e8 --- /dev/null +++ b/unittests/physics/LR-test.hpp @@ -0,0 +1,140 @@ +#pragma once + +#include "RI/physics/LR.h" +#include "Test_Helpers.hpp" + +namespace LR_Test +{ + +using namespace Test_Helpers; +static const std::string ref_path = "unittests/physics/LR-ref"; + +// -------- main test -------- +static void main(int argc, char *argv[]) +{ + int mpi_init_provide; + MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &mpi_init_provide); + + // ====== config: 1 atom, 3x1x1 k-grid, nocc=1, nvirt=2, nabf=2, nw=3 ====== + const TA atom0 = 0; + const int nk = 3; + const int nocc = 1, nvirt = 2, nband = nocc + nvirt; + const int nabf = 2, nw = 3; + + // 3x1x1 uniform k-grid fractional coordinates + std::vector kfrac(nk); + for (int i = 0; i < nk; ++i) + kfrac[i] = {{double(i)/3.0, 0.0, 0.0}}; + + // Cs: shape {nabf=2, nw=3, nw=3}, values vary with R + std::map>> Cs; + for (auto R : {std::array{0,0,0}, std::array{1,0,0}, std::array{-1,0,0}}) { + const double w = (R[0]==0 && R[1]==0 && R[2]==0) ? 1.0 : 0.5 / (std::abs(R[0])+std::abs(R[1])+std::abs(R[2])); + Cs[atom0][{atom0, R}] = RI::Tensor({nabf, nw, nw}); + for (int mu = 0; mu < nabf; ++mu) + for (int s = 0; s < nw; ++s) + for (int t = 0; t < nw; ++t) + Cs[atom0][{atom0, R}](mu,s,t) = Tdata(w * (mu+1) * (s+1) * (t+1), 0.0); + } + + // Vs / Ws: shape {nabf=2, nabf=2}, values vary with R + auto make_V = [&]() { + std::map>> V; + for (auto R : {std::array{0,0,0}, std::array{1,0,0}, std::array{-1,0,0}}) { + const double w = (R[0]==0 && R[1]==0 && R[2]==0) ? 0.01 : 0.005 / (std::abs(R[0])+std::abs(R[1])+std::abs(R[2])); + V[atom0][{atom0, R}] = RI::Tensor({nabf, nabf}); + for (int i = 0; i < nabf; ++i) + for (int j = 0; j < nabf; ++j) + V[atom0][{atom0, R}](i,j) = (i==j) ? Tdata(w, 0.0) : Tdata(w*0.5, 0.0); + } + return V; + }; + auto Vs = make_V(); + auto Ws = make_V(); + + // ====== init LR with MPI distribution ====== + RI::LR lr; + lr.init(kfrac, nocc, nvirt); + lr.set_parallel(MPI_COMM_WORLD, 1, nk, {{1,1,1}}); + lr.set_Cs(Cs, 0.0, {atom0}, {atom0}); + lr.set_Vs(Vs, 0.0, {atom0}, {atom0}); + lr.set_Ws(Ws, 0.0, {atom0}, {atom0}); + + // build q_list / q2kpair from distributed k1, k2 + { + auto mod1 = [](double x) { double r = x - std::floor(x); return r < 0 ? r+1.0 : r; }; + for (int k1 : lr.k1_indices) + for (int k2 : lr.k2_indices) { + Tk q = {{mod1(kfrac[k2][0]-kfrac[k1][0]), + mod1(kfrac[k2][1]-kfrac[k1][1]), + mod1(kfrac[k2][2]-kfrac[k1][2])}}; + lr.q2kpair[q].emplace_back(k1, k2); + } + for (auto& p : lr.q2kpair) lr.q_list.push_back(p.first); + } + + // map_psi: only for distributed k_indices + for (int ik : lr.k_indices) { + lr.map_psi[ik][atom0] = RI::Tensor({nband, nw}); + for (int m = 0; m < nband; ++m) + for (int t = 0; t < nw; ++t) + lr.map_psi[ik][atom0](m,t) + = Tdata(0.1 * (m+1.0)*(t+1.0)*(ik+1.0), + 0.1 * (m+0.5)*(t+0.5)*(ik+0.5)); + } + + // compute all 5 results + std::ofstream ofs("/dev/null"); + lr.cal_Csk_ao_mo("Cs_", ofs); + auto V_A = lr.cal_cvc_mo_k_hartree_onthefly({"O","V","O","V"}, "Vs_", true); + auto V_B = lr.cal_cvc_mo_k_hartree_onthefly({"O","V","O","V"}, "Vs_", false); + auto W_A = lr.cal_cvc_mo_k_onthefly({"O","O","V","V"}, "Ws_", true); + auto W_B = lr.cal_cvc_mo_k_onthefly({"V","O","O","V"}, "Ws_", false); + + // ====== gather per-rank results to rank 0, then regression ====== + int nproc; MPI_Comm_size(MPI_COMM_WORLD, &nproc); + auto Csk_gathered = gather_map(lr.Csk_ao_mo, 0, MPI_COMM_WORLD, 0); + auto VA_gathered = gather_map(V_A, 0, MPI_COMM_WORLD, 1); + auto VB_gathered = gather_map(V_B, 0, MPI_COMM_WORLD, 2); + auto WA_gathered = gather_map(W_A, 0, MPI_COMM_WORLD, 3); + auto WB_gathered = gather_map(W_B, 0, MPI_COMM_WORLD, 4); + + if (nproc > 1) MPI_Barrier(MPI_COMM_WORLD); + + if (argc >= 2 && std::string(argv[1]) == "--write-ref") + { + assert(nproc == 1 && "--write-ref requires exactly 1 MPI process"); + std::ofstream ofs_ref(ref_path); + assert(ofs_ref && "Failed to open ref file for writing"); + ofs_ref << std::scientific << std::setprecision(REF_PRECISION); + write_map(ofs_ref, Csk_gathered); + write_map(ofs_ref, VA_gathered); + write_map(ofs_ref, VB_gathered); + write_map(ofs_ref, WA_gathered); + write_map(ofs_ref, WB_gathered); + std::cout << "LR_ref written to " << ref_path << std::endl; + } + else + { + int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank); + if (rank == 0) + { + std::ifstream ifs_ref(ref_path); + assert(ifs_ref && "LR_ref not found; run with --write-ref to generate"); + TensorMap ref; + read_map(ifs_ref, ref); compare_map(Csk_gathered, ref); + read_map(ifs_ref, ref); compare_map(VA_gathered, ref); + read_map(ifs_ref, ref); compare_map(VB_gathered, ref); + read_map(ifs_ref, ref); compare_map(WA_gathered, ref); + read_map(ifs_ref, ref); compare_map(WB_gathered, ref); + } + } + + lr.free_Cs(); + lr.free_Vs(); + lr.free_Ws(); + + MPI_Finalize(); +} + +} diff --git a/unittests/physics/Test_Helpers.hpp b/unittests/physics/Test_Helpers.hpp new file mode 100644 index 0000000..67f63a8 --- /dev/null +++ b/unittests/physics/Test_Helpers.hpp @@ -0,0 +1,178 @@ +// =================== +// Author: Ziqing Guan +// date: 2026.08.02 +// =================== + +#pragma once + +#include "RI/global/Tensor.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REF_PRECISION 9 // significant digits when writing reference file +#define REF_TOLERANCE 1e-7 // tolerance when comparing against reference + +namespace Test_Helpers +{ + +using Tdata = std::complex; +using TA = int; +using Tcell = int; +constexpr int Ndim = 3; +using TAC = std::pair>; +using Tk = std::array; + +// ====== type aliases ====== +using TensorMap = std::map>>; +using HartreeMap = std::map>>>; + +// ====== write (recursive dispatch: map -> recurse, Tensor -> leaf) ====== + +// leaf: Tensor +template +inline void write_map(std::ostream &os, const RI::Tensor &t) +{ + os << t.shape.size(); + for (auto d : t.shape) os << " " << d; + std::size_t total = 1; + for (auto d : t.shape) total *= d; + for (std::size_t i = 0; i < total; ++i) + os << " " << std::setprecision(REF_PRECISION) << std::scientific + << t.ptr()[i].real() << " " << t.ptr()[i].imag(); + os << "\n"; +} + +// recurse: map +template +inline void write_map(std::ostream &os, const std::map &m) +{ + os << m.size() << "\n"; + for (auto &[k, v] : m) { + os << k << " "; + write_map(os, v); + } +} + +// ====== read ====== + +// leaf: Tensor +template +inline void read_map(std::istream &is, RI::Tensor &t) +{ + std::size_t ndim; is >> ndim; + std::vector shape(ndim); + std::size_t total = 1; + for (std::size_t d = 0; d < ndim; ++d) { + is >> shape[d]; total *= shape[d]; + } + t = RI::Tensor(shape); + for (std::size_t i = 0; i < total; ++i) { + double re, im; is >> re >> im; t.ptr()[i] = T(re, im); + } +} + +// recurse: map +template +inline void read_map(std::istream &is, std::map &m) +{ + m.clear(); + std::size_t n; is >> n; + for (std::size_t i = 0; i < n; ++i) { + K key; is >> key; + read_map(is, m[key]); + } +} + +// ====== compare ====== + +// leaf: Tensor +template +inline void compare_map(const RI::Tensor &a, const RI::Tensor &b) +{ + assert(a.shape.size() == b.shape.size()); + std::size_t total = 1; + for (std::size_t d = 0; d < a.shape.size(); ++d) { + assert(a.shape[d] == b.shape[d]); + total *= a.shape[d]; + } + for (std::size_t i = 0; i < total; ++i) + assert(std::abs(a.ptr()[i] - b.ptr()[i]) < REF_TOLERANCE); +} + +// recurse: map +template +inline void compare_map(const std::map &a, const std::map &b) +{ + assert(a.size() == b.size()); + for (auto &[k, va] : a) { + assert(b.count(k)); + compare_map(va, b.at(k)); + } +} + +// ====== merge (for MPI gather) ====== + +// base: non-map value — move individual entries +template +inline void merge_into(std::map &target, std::map &&source) +{ + for (auto &[k, v] : source) + target[k] = std::move(v); +} + +// recurse: value is another map — merge inner maps element by element +template +inline void merge_into(std::map> &target, + std::map> &&source) +{ + for (auto &[k, inner] : source) + merge_into(target[k], std::move(inner)); +} + +// ====== MPI gather ====== + +template +inline M gather_map(const M &local, int root, MPI_Comm comm, int tag = 0) +{ + int rank, nproc; + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &nproc); + + std::ostringstream oss; + write_map(oss, local); + std::string data = oss.str(); + int len = static_cast(data.size()); + + if (rank == root) + { + M combined = local; + for (int r = 0; r < nproc; ++r) + { + if (r == root) continue; + int remote_len; + MPI_Recv(&remote_len, 1, MPI_INT, r, tag, comm, MPI_STATUS_IGNORE); + std::string remote_data(remote_len, '\0'); + MPI_Recv(&remote_data[0], remote_len, MPI_CHAR, r, tag, comm, MPI_STATUS_IGNORE); + std::istringstream iss(remote_data); + M remote; + read_map(iss, remote); + merge_into(combined, std::move(remote)); + } + return combined; + } + else + { + MPI_Send(&len, 1, MPI_INT, root, tag, comm); + MPI_Send(data.data(), len, MPI_CHAR, root, tag, comm); + return {}; + } +} + +} diff --git a/unittests/ri/Cell_Nearest-test.hpp b/unittests/ri/Cell_Nearest-test.hpp index b0f94c6..7a23e29 100644 --- a/unittests/ri/Cell_Nearest-test.hpp +++ b/unittests/ri/Cell_Nearest-test.hpp @@ -150,5 +150,112 @@ namespace Cell_Nearest_Test 2 1 | 0 -1 */ } + + test_direction(); + } + + static void test_direction() + { + using TA = std::string; + using Tcell = int; + constexpr int Ndim = 3; + using Tpos = double; + constexpr int Npos = 3; + using TC = std::array; + using Tatom_pos = std::array; + + std::map atoms_pos; + std::array latvec; + std::array period; + + // graphene-like hexagonal lattice, period=5 in a,b directions + latvec[0] = {10, 0, 0}; + latvec[1] = {5, std::sqrt(75.0), 0}; + latvec[2] = {0, 0, 100}; + period = {5, 5, 1}; + + // both atoms at origin -> Ryx = {0,0,0} + atoms_pos["H"] = {0, 0, 0}; + atoms_pos["C"] = {0, 0, 0}; + + RI::Cell_Nearest cn; + cn.init(atoms_pos, latvec, period); + std::cout< switch to (3,-1,0) + { + run("Test C: cell( 3, 4,0)", "H", "C", {3,4,0}); + // Expected: (3,-1,0) dist ~ 26.4575 + } + + // Test D: cell(4,3,0) + // 2 cells tie at dist = 10*sqrt(7): + // (-1,-2,0) (a=-1,b=-1) and (-1,3,0) (a=-1,b=0) + // |b|=0 < |b|=1 -> switch to (-1,3,0) + { + run("Test D: cell( 4, 3,0)", "H", "C", {4,3,0}); + // Expected: (-1,3,0) dist ~ 26.4575 + } + + // Test E: cell(6,7,0) + // 2 cells tie at dist = 10*sqrt(7): + // (1,-3,0) (a=-1,b=-2) and (1,2,0) (a=-1,b=-1) + // |b|=1 < |b|=2 -> switch to (1,2,0) + { + run("Test E: cell( 6, 7,0)", "H", "C", {6,7,0}); + // Expected: (1,2,0) dist ~ 26.4575 + } + + // Test F: cell(7,6,0) + // 2 cells tie at dist = 10*sqrt(7): + // (-3,1,0) (a=-2,b=-1) and (2,1,0) (a=-1,b=-1) + // |b| both=1, then |a|=1 < |a|=2 -> switch to (2,1,0) + { + run("Test F: cell( 7, 6,0)", "H", "C", {7,6,0}); + // Expected: (2,1,0) dist ~ 26.4575 + } + /* + C C 0 0 0 + H 0 0 0 + H C 0 0 0 + H 0 0 0 + + Test A: cell( 2, 2,0): cell_nearest=(-3,2,0) dist=26.4575 + Test B: cell(-2,-2,0): cell_nearest=(3,-2,0) dist=26.4575 + Test C: cell( 3, 4,0): cell_nearest=(3,-1,0) dist=26.4575 + Test D: cell( 4, 3,0): cell_nearest=(-1,3,0) dist=26.4575 + Test E: cell( 6, 7,0): cell_nearest=(1,2,0) dist=26.4575 + Test F: cell( 7, 6,0): cell_nearest=(2,1,0) dist=26.4575 + */ } } \ No newline at end of file