diff --git a/include/RI/distribute/Distribute_Equally.h b/include/RI/distribute/Distribute_Equally.h index 3f1acca..60d03bc 100644 --- a/include/RI/distribute/Distribute_Equally.h +++ b/include/RI/distribute/Distribute_Equally.h @@ -56,6 +56,17 @@ 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); } } 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..a405624 100644 --- a/include/RI/global/Global_Func-1.h +++ b/include/RI/global/Global_Func-1.h @@ -83,6 +83,34 @@ namespace Global_Func return find( ptr->second, keys... ); } + // These helper functions are used to replace Global_Func::find, + // since the target is not Tensor, but another map. + template + static inline const typename Map::mapped_type find_map(const Map& map, const Key& key) + { + const auto& it = map.find(key); + if (it != map.end()) return it->second; + return typename Map::mapped_type(); + } + + template + static inline const typename Map::mapped_type::mapped_type find_map( + const Map& map, const Key1& key1, const Key2& key2) + { + using InnerMap = typename Map::mapped_type; + using Target = typename InnerMap::mapped_type; + + const auto it1 = map.find(key1); + if (it1 != map.end()) + { + const auto& inner = it1->second; + const auto it2 = inner.find(key2); + if (it2 != inner.end()) + return it2->second; + } + return Target(); + } + // in_set(3, {2,3,5,7}) // Peize Lin add 2022.05.26 template 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..dbb3f7f 100644 --- a/include/RI/global/Shape_Vector.h +++ b/include/RI/global/Shape_Vector.h @@ -29,6 +29,18 @@ 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]; + } + Shape_Vector(std::vector&& 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_; } diff --git a/include/RI/global/Tensor.h b/include/RI/global/Tensor.h index 4f5a2c9..a106726 100644 --- a/include/RI/global/Tensor.h +++ b/include/RI/global/Tensor.h @@ -45,6 +45,15 @@ class Tensor Tensor transpose() const; Tensor dagger() const; + Tensor conjugate() const; + + // get maximum absolute value among all elements + Global_Func::To_Real_t max_abs() const; + + /// permute from the input index order to {0, 1, 2, ..., N} + Tensor permute_from(const std::vector& order) const; + /// permute from {0, 1, 2, ..., N } to the input new index order + // Tensor permute_to(const std::vector) 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..027d0f1 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::max_abs() const +{ + using T_res = Global_Func::To_Real_t; + T_res max_val = 0.0; + std::size_t size = this->get_shape_all(); + for (size_t i = 0; i < size; ++i) + { + T_res abs_val = std::abs((*this->data)[i]); + if (abs_val > max_val) + { + max_val = abs_val; + } + } + return max_val; +} + template Global_Func::To_Real_t Tensor::norm(const double p) const { @@ -356,4 +391,74 @@ std::array,N2>,N1>,N0> to_array(const Ten return a; } +// check whether an array is a permutation of {from, from+1, ..., to} +inline bool is_permutation(const std::vector& arr, const std::size_t from, const std::size_t to) +{ + std::vector seen(arr.size(), false); + if (arr.size() != to - from + 1) return false; + for (std::size_t i = 0; i < arr.size(); ++i) { + if (arr[i] < from || arr[i] > to) return false; + const std::size_t idx = arr[i] - from; + if (seen[idx]) return false; + seen[idx] = true; + } + return true; +} + +template +Tensor Tensor::permute_from(const std::vector& order) const +{ + const std::size_t N = order.size(); + assert(this->shape.size() == N); + assert(is_permutation(order, 0, N-1)); + Shape_Vector new_shape(this->shape); + for (std::size_t d = 0; d < N; ++d) + new_shape[d] = this->shape[order[d]]; + + Tensor t(new_shape); + + /* N-permutation + for (std::size_t i0 = 0; i0 < new_shape[0], ++i0) + for (std::size_t i1 = 0, i1 < new_shape[1], ++i1) + ... + for (std::size_t iN = 0, iN < new_shape[N], ++iN) + t(i0, i1, ..., iN) = this->operator()(order(i0), order(i1), ..., order(N)) + */ + std::vector strides(N); + strides[N - 1] = 1; + for (int d = N - 2; d >= 0; --d) + strides[d] = strides[d + 1] * this->shape[d + 1]; + + std::vector new_strides(N); + new_strides[N - 1] = 1; + for (int d = N - 2; d >= 0; --d) + new_strides[d] = new_strides[d + 1] * new_shape[d + 1]; + + + std::vector new_idx(N, 0); + std::vector idx(N, 0); + for (std::size_t i = 0; i < t.get_shape_all(); ++i) // new 1D-index + { + // new ND-index + std::size_t tmp = i; + for (std::size_t d = 0; d < N; ++d) + { + new_idx[d] = tmp / new_strides[d]; + tmp %= new_strides[d]; + } + // old ND-index + for (std::size_t d = 0; d < N; ++d) + idx[order[d]] = new_idx[d]; + + // old 1D-index + std::size_t i0 = 0; + for (std::size_t d = 0; d < N; ++d) + i0 += idx[d] * strides[d]; + + (*t.data)[i] = (*this->data)[i0]; + } + + return t; +} + } \ No newline at end of file diff --git a/include/RI/global/Tensor_Multiply-32.hpp b/include/RI/global/Tensor_Multiply-32.hpp index a46c3f9..9c8cdca 100644 --- a/include/RI/global/Tensor_Multiply-32.hpp +++ b/include/RI/global/Tensor_Multiply-32.hpp @@ -80,6 +80,49 @@ namespace Tensor_Multiply return Txy; } + // 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(), ny, Vy.ptr(), 1, + Tdata(0.0), Txy.ptr(), 1); + 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(), Txy.get_shape_all(), Vy.ptr(), 1, + Tdata(0.0), Txy.ptr(), 1); + return Txy; + } } } diff --git a/include/RI/physics/Exx.h b/include/RI/physics/Exx.h index e7457d1..bdb8ee9 100644 --- a/include/RI/physics/Exx.h +++ b/include/RI/physics/Exx.h @@ -29,6 +29,10 @@ class Exx constexpr static std::size_t Npos = Ndim; // tmp using Tatom_pos = std::array; // tmp + Exx(const std::string method = "loop3") { this->method = method; } + Exx(Exx&&) noexcept = default; + Exx& operator=(Exx&&) noexcept = default; + void set_parallel( const MPI_Comm &mpi_comm, const std::map &atoms_pos, @@ -96,6 +100,7 @@ class Exx void free_dVs(const std::string &save_name_suffix=""); void free_dCRs(const std::string &save_name_suffix=""); void free_dVRs(const std::string &save_name_suffix=""); + void free_cvc() { this->cvc_.clear(); } public: LRI lri; @@ -122,6 +127,9 @@ class Exx }; Flag_Save_Result flag_save_result; + std::string method = "loop3"; // "loop3", "cvc" + std::map>>>> cvc_; + std::map>>>>& get_cvc() const { return this->cvc_; } }; } diff --git a/include/RI/physics/Exx.hpp b/include/RI/physics/Exx.hpp index 2bb8be3..a1073eb 100644 --- a/include/RI/physics/Exx.hpp +++ b/include/RI/physics/Exx.hpp @@ -266,12 +266,21 @@ void Exx::cal_Hs( if(!this->flag_finish.Ds_delta) this->Hs.clear(); - this->lri.cal_loop3( - {Label::ab_ab::a0b0_a1b1, - Label::ab_ab::a0b0_a1b2, - Label::ab_ab::a0b0_a2b1, - Label::ab_ab::a0b0_a2b2}, - this->Hs); + if (this->method == "loop3") + this->lri.cal_loop3( + { Label::ab_ab::a0b0_a1b1, + Label::ab_ab::a0b0_a1b2, + Label::ab_ab::a0b0_a2b1, + Label::ab_ab::a0b0_a2b2 }, + this->Hs); + else if (this->method == "cvc") + { + if (this->cvc_.empty()) + this->cvc_ = this->lri.cal_cvc(); + this->Hs = this->lri.constract_cvc_ds(this->cvc_); + } + else + assert(false && "Unknown method in Exx::cal_Hs"); //if() this->energy = this->post_2D.cal_energy( diff --git a/include/RI/physics/Hartree.h b/include/RI/physics/Hartree.h new file mode 100644 index 0000000..2449380 --- /dev/null +++ b/include/RI/physics/Hartree.h @@ -0,0 +1,100 @@ +// =================== +// Author: Ziqing Guan +// date: 2025.12.26 +// =================== + +#pragma once +#include "../global/Global_Func-2.h" +#include "../global/Tensor.h" +#include "../ri/LRI.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 Tatom_pos = std::array; // tmp + + void set_parallel( + const MPI_Comm &mpi_comm, + const std::map &atoms_pos, + const std::array &latvec, + const std::array &period) + { + this->lri.set_parallel( + mpi_comm, atoms_pos, latvec, period, + {}/*no label*/); + 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->lri.mpi_comm, std::move(Cs), listI, listJ); + this->lri.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->lri.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->lri.mpi_comm, std::move(Vs), listI, listJ); + this->lri.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->lri.free_tensors_map2("Vs_"+save_name_suffix); + this->flag_finish.Vs = false; + }; + + + +public: + LRI lri; + + struct Flag_Finish + { + bool stru=false; + bool Cs=false; + bool Vs=false; + bool Ds=false; + }; + Flag_Finish flag_finish; + +}; + +} \ No newline at end of file diff --git a/include/RI/physics/LR.h b/include/RI/physics/LR.h new file mode 100644 index 0000000..e476311 --- /dev/null +++ b/include/RI/physics/LR.h @@ -0,0 +1,147 @@ +#pragma once +#include "./Exx.h" + +namespace RI +{ +// Difference from Exx: +// 1. in the `cal_force` function, two density matrices can be different +// 2. set_Cs and set_Vs can parallelize the input tensors according to listI and listJ +// 3. add set_Ws and free_Ws to setup screened Coulomb interaction +// 4. feat: cvc_mo in k space +template +class LR : public Exx +{ +public: + using TC = std::array; + using TAC = std::pair; + using Tdata_real = Global_Func::To_Real_t; + using Tk = std::array; + + LR(const std::string method = "cvc") { this->method = method; } // Exx default mehtod: "loop3" + LR(Exx&& exx) : Exx::Exx(std::move(exx)) {} + using Exx::cal_force; + // New function: cal_force with two different density matrices + void cal_force(const std::map>, Tensor>>& Ds_left, + const std::array& save_names_suffix = { "","","","","" }) // "Cs","Vs","Ds","dCs","dVs" + { + // The only difference from Exx::cal_force: save Ds_left if not empty + if (!Ds_left.empty()) + this->post_2D.saves["Ds_" + save_names_suffix[2]] = this->post_2D.set_tensors_map2(Ds_left); + this->cal_force(save_names_suffix); + } + + 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->lri.mpi_comm, std::move(Cs), listI, listJ); + this->lri.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 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->lri.mpi_comm, std::move(Vs), listI, listJ); + this->lri.set_tensors_map2( + Vs, + {Label::ab::a0b0}, + {{"threshold_filter", threshold}, {"flag_comm", false}}, + "Vs_"+save_name_suffix ); + this->flag_finish.Vs = true; + }; + + // setup screened Coulomb interaction and parallelize according to listI and listJ + 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->lri.mpi_comm, std::move(Ws), listI, listJ); + this->lri.set_tensors_map2( + Ws, + {Label::ab::a0b0}, + {{"threshold_filter", threshold}, {"flag_comm", false}}, + "Ws_"+save_name_suffix ); + this->flag_Ws = true; + }; + + void free_Ws(const std::string &save_name_suffix="") + { + this->lri.free_tensors_map2("Ws_"+save_name_suffix); + this->flag_Ws = false; + }; + + bool flag_Ws = false; + + std::vector k1_indices; + std::vector k2_indices; + std::vector k_indices; // k1∪k2 (for psi/mo transform) + std::vector kindex_map; // index → Tk fractional coord + std::vector list_I; + std::vector list_J; + std::vector list_IJ; // I∪J + std::size_t nocc = 0; + std::size_t nvirt = 0; + + std::vector q_list; + std::map>> q2kpair; // q → list of (k1_idx, k2_idx) + + std::map>> map_psi; //> + std::map>> Csk_ao_mo; + + /// @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->lri.data_pool.at(save_name).Ds_ab; + this->Csk_ao_mo = this->lri.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->lri.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->lri.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); + } +}; + +} 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-cal_cvc.hpp b/include/RI/ri/LRI-cal_cvc.hpp new file mode 100644 index 0000000..2b6978b --- /dev/null +++ b/include/RI/ri/LRI-cal_cvc.hpp @@ -0,0 +1,310 @@ +#pragma once + +#include "LRI.h" +#include "LRI_Cal_Aux.h" +#include "../global/Array_Operator.h" +#include "../global/Tensor_Multiply.h" +#include +#include +#ifdef __MKL_RI +#include +#endif + +namespace RI +{ + +template +std::map>, std::map>, Tensor>>>> +LRI::cal_cvc() +{ + using namespace Array_Operator; + + const Data_Pack_Wrapper data_wrapper(this->data_pool, this->data_ab_name); + const LRI_Cal_Tools tools(this->period, this->data_pool, this->data_ab_name); + + #ifdef __MKL_RI + const std::size_t mkl_threads = mkl_get_max_threads(); + mkl_set_num_threads(1); + #endif + + std::map>>>> cvc; + /// ??? + std::map lock_cvc_result_add_map = LRI_Cal_Aux::init_lock_result({ Label::ab_ab::a0b0_a2b2 }, this->parallel->list_A, cvc); + + #pragma omp parallel + { + std::map>>>> cvc_thread; +/* pseudo code +for I + for J + for K + [CV]_{IK,J}=C^I_{IK}V_{IJ} + for L + [CVC]_{IJKL}+=[CV]_{IK,J} C^J_{JL} #1 + for L + for K + [CV]_{IK,L}=C^I_{IK}V_{IL} + for J + [CVC]_{IJKL}+=[CV]_{IK,L} C^L_{LJ} #2 +for K + for J + for I + [CV]_{KI,J}=C^K_{KI}V_{KJ} + for L + [CVC]_{IJKL}+=[CV]_{KI,J}C^J_{JL} #3 + for L + for I + [CV]_{KI,L}=C^K_{KI}V_{KL} + for J + [CVC]_{IJKL}+=[CV]_{KI,L}C^L_{LJ} #4 +*/ + // Main Problems: + // How to filter list_I, list_Aa02, list_K, list_L only using the range of Vs? + // How to deal with thread lock? + // 转置无法通过改变乘法顺序避免： 输出KL连续，输入KL分别在两个C上 + // for debug + auto print_a = [](const std::vector& vec, const std::string name) -> void + { + std::cout << name << ": "; + for (auto& v : vec) { std::cout << v << " "; } + std::cout << std::endl; + }; + auto print_ac = [](const std::vector& vec, const std::string name) -> void + { + std::cout << name << ": "; + for (auto& v : vec) { std::cout << v.first << " "; } + std::cout << std::endl; + }; + const std::vector list_I0 = LRI_Cal_Aux::filter_list_map(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).a01, data_wrapper(Label::ab::a).Ds_ab); + const std::vector list_J0 = LRI_Cal_Aux::filter_list_map(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).b01, data_wrapper(Label::ab::b).Ds_ab); + const std::vector list_K = LRI_Cal_Aux::filter_list_set(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).a2, data_wrapper(Label::ab::a).index_Ds_ab[0]); + const std::vector list_L = LRI_Cal_Aux::filter_list_set(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).b2, data_wrapper(Label::ab::b).index_Ds_ab[0]); + // filter Cs by the range of Vs + const std::vector list_I = LRI_Cal_Aux::filter_list_map(list_I0, data_wrapper(Label::ab::a0b0).Ds_ab); + const std::vector list_J = LRI_Cal_Aux::filter_list_set(list_J0, data_wrapper(Label::ab::a0b0).index_Ds_ab[0]); + // print_a(list_I0, "list_I0"); + // print_ac(list_J0, "list_J0"); + // print_ac(list_K, "list_K"); + // print_ac(list_L, "list_L"); + // print_a(list_I, "list_I"); + // print_ac(list_J, "list_J"); + + // allocate + // for (TA I : list_I) + // for (TAC J : list_J) + // for (TAC K : list_K) + // { + // const Tensor& C_I_IK = tools.get_Ds_ab(Label::ab::a, I, K); + // const int& ni = C_I_IK.shape[1]; + // const int& nk = C_I_IK.shape[2]; + // for (TAC L : list_L) + // { + // const TC R_KL = (L.second - K.second) % period; + // const Tensor& C_J_JL = tools.get_Ds_ab(Label::ab::b, J, L); + // const int& nj = C_J_JL.shape[1]; + // const int& nl = C_J_JL.shape[2]; + // cvc_thread[I][J][K.first][{L.first, R_KL}] = Tensor({ ni, nj, nk, nl }); + // } + // } + + // calculate + for (TA I : list_I) + { + if (this->filter_atom->filter_for1(Label::ab_ab::a0b0_a2b2, I)) continue; // restrict I in the irreducible sector + auto& cvc_thread_I = cvc_thread[I]; + #pragma omp for schedule(static) nowait + for (TAC J : list_J) // term 1 + { + if (this->filter_atom->filter_for32(Label::ab_ab::a0b0_a2b2, I, J, J)) continue; // restrict (I, J) in the irreducible sector + const Tensor& V_IJ = tools.get_Ds_ab(Label::ab::a0b0, I, J); + if (V_IJ.empty()) continue; + // symmetry: check if (I, J) in irreducible sector + auto& cvc_thread_IJ = cvc_thread_I[J]; + for (TAC K : list_K) + { + const Tensor& C_I_IK = tools.get_Ds_ab(Label::ab::a, I, K); + if (C_I_IK.empty()) continue; + // CV_{IK,J}=C^I_{IK}V_{IJ} + const Tensor CV_IK_J = Tensor_Multiply::x1x2y1_ax1x2_ay1(C_I_IK, V_IJ); + auto& cvc_thread_IJK = cvc_thread_IJ[K.first]; + for (TAC L : list_L) + { + const TC R_KL = (L.second - K.second) % period; + // std::array key = { I, J, K, L }; + // [CVC]_{IJKL}+=[CV]_{IK,J}C^J_{JL} + const Tensor& C_J_JL = tools.get_Ds_ab(Label::ab::b, J, L); + if (C_J_JL.empty()) continue; + //(ika) * (ajl) = (ikjl) -> (ijkl) + // cvc_thread[{I, J, K, L}][R_KL] += einsum("ika, ajl -> ijkl", CV_IK_J, C_J_JL); + // cvc_thread[I][J][K.first][{L.first, R_KL}] += Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_IK_J, C_J_JL).permute_from({ 0,2,1,3 }); + LRI_Cal_Aux::add_Ds( + Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_IK_J, C_J_JL).permute_from({ 0,2,1,3 }), + cvc_thread_IJK[{L.first, R_KL}]); + } + } + LRI_Cal_Aux::add_Ds_omp_try_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + } + } + LRI_Cal_Aux::add_Ds_omp_wait_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + + for (TA I : list_I) + { + if (this->filter_atom->filter_for1(Label::ab_ab::a0b0_a2b2, I)) continue; + auto& cvc_thread_I = cvc_thread[I]; + #pragma omp for schedule(static) nowait + for (TAC L : list_L) // term 2 + { + const Tensor& V_IL = tools.get_Ds_ab(Label::ab::a0b0, I, L); + if (V_IL.empty()) continue; + for (TAC K : list_K) + { + const TC& R_KL = (L.second - K.second) % period; + const Tensor& C_I_IK = tools.get_Ds_ab(Label::ab::a, I, K); + if (C_I_IK.empty()) continue; + // CV_{IK,L}: a1a2b0 = a0a1a2 * a0b0 + const Tensor CV_IK_L = Tensor_Multiply::x1x2y1_ax1x2_ay1(C_I_IK, V_IL); + for (TAC J : list_J) + { + if (this->filter_atom->filter_for32(Label::ab_ab::a0b0_a2b2, I, J, J)) continue; // restrict (I, J) in the irreducible sector + // [CVC]_{IJKL}+=[CV]_{IK,L}C^L_{LJ} + const Tensor& C_L_LJ = tools.get_Ds_ab(Label::ab::b, L, J); + if (C_L_LJ.empty()) continue; + // (ika) * (alj) = (iklj) -> (ijkl) + // cvc_thread[{I, J, K, L}][R_KL] += einsum("ika, alj -> ijkl", CV_IK_L, C_L_LJ); + // cvc_thread[I][J][K.first][{L.first, R_KL}] += Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_IK_L, C_L_LJ).permute_from({ 0,3,1,2 }); + LRI_Cal_Aux::add_Ds( + Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_IK_L, C_L_LJ).permute_from({ 0,3,1,2 }), + cvc_thread_I[J][K.first][{L.first, R_KL}]); + } + } + LRI_Cal_Aux::add_Ds_omp_try_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + } + } + LRI_Cal_Aux::add_Ds_omp_wait_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + + for (TAC K : list_K) + { + #pragma omp for schedule(static) nowait + for (TAC J : list_J) //term 3 + { + if (this->filter_atom->filter_for1(Label::ab_ab::a0b0_a1b2, J)) continue; // restrict J in the irreducible sector + const Tensor& V_KJ = tools.get_Ds_ab(Label::ab::a0b0, K, J); + if (V_KJ.empty()) continue; + for (TA I : list_I) + { + if (this->filter_atom->filter_for32(Label::ab_ab::a0b0_a2b2, I, J, J)) continue; // restrict (I, J) in the irreducible sector + const Tensor& C_K_KI = tools.get_Ds_ab(Label::ab::a, K, I); + if (C_K_KI.empty()) continue; + //[CV]_{KI,J}=C^K_{KI}V_{KJ} + const Tensor CV_KI_J = Tensor_Multiply::x1x2y1_ax1x2_ay1(C_K_KI, V_KJ); + auto& cvc_thread_IJK = cvc_thread[I][J][K.first]; + for (TAC L : list_L) + { + const TC& R_KL = (L.second - K.second) % period; + // [CVC]_{IJKL}+=[CV]_{KI,J}C^J_{JL} #(3) + const Tensor& C_J_JL = tools.get_Ds_ab(Label::ab::b, J, L); + if (C_J_JL.empty()) continue; + // (kia) * (ajl) = (kijl)->(ijkl) + // cvc_thread[{I, J, K, L}][R_KL] += einsum("kia, ajl -> ijkl", CV_KI_J, C_J_JL); + // cvc_thread[I][J][K.first][{L.first, R_KL}] += Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_KI_J, C_J_JL).permute_from({ 1,2,0,3 }); + LRI_Cal_Aux::add_Ds( + Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_KI_J, C_J_JL).permute_from({ 1,2,0,3 }), + cvc_thread_IJK[{L.first, R_KL}]); + } + } + LRI_Cal_Aux::add_Ds_omp_try_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + } + } + LRI_Cal_Aux::add_Ds_omp_wait_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + + for (TAC K : list_K) + { + #pragma omp for schedule(static) nowait + for (TAC L : list_L) // term 4 + { + const Tensor& V_KL = tools.get_Ds_ab(Label::ab::a0b0, K, L); + if (V_KL.empty()) continue; + const TC& R_KL = (L.second - K.second) % period; + for (TA I : list_I) + { + if (this->filter_atom->filter_for1(Label::ab_ab::a0b0_a2b2, I)) continue; // restrict I in the irreducible sector + const Tensor& C_K_KI = tools.get_Ds_ab(Label::ab::a, K, I); + if (C_K_KI.empty()) continue; + // CV_{IK,L}: a1a2b0 = a0a1a2 * a0b0 + const Tensor CV_KI_L = Tensor_Multiply::x1x2y1_ax1x2_ay1(C_K_KI, V_KL); + for (TAC J : list_J) + { + if (this->filter_atom->filter_for32(Label::ab_ab::a0b0_a2b2, I, J, J)) continue; // restrict (I, J) in the irreducible sector + // [CVC]_{IJKL}+=[CV]_{IK,L}C^L_{LJ} + const Tensor& C_L_LJ = tools.get_Ds_ab(Label::ab::b, L, J); + if (C_L_LJ.empty()) continue; + // (kia) * (alj) = (kilj) -> (ijkl) + // cvc_thread[{I, J, K, L}][R_KL] += einsum("kia, alj -> ijkl", CV_KI_L, C_L_LJ); + // cvc_thread[I][J][K.first][{L.first, R_KL}] += Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_KI_L, C_L_LJ).permute_from({ 1,3,0,2 }); + LRI_Cal_Aux::add_Ds( + Tensor_Multiply::x0x1y1y2_x0x1a_ay1y2(CV_KI_L, C_L_LJ).permute_from({ 1,3,0,2 }), + cvc_thread[I][J][K.first][{L.first, R_KL}]); + } + } + LRI_Cal_Aux::add_Ds_omp_try_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + } + } + LRI_Cal_Aux::add_Ds_omp_wait_map(cvc_thread, cvc, lock_cvc_result_add_map, 1.0); + } // end #pragma omp parallel + + LRI_Cal_Aux::destroy_lock_result(lock_cvc_result_add_map, cvc); + + #ifdef __MKL_RI + mkl_set_num_threads(mkl_threads); + #endif + + malloc_trim(0); + return cvc; +} // end LRI::cal_cvc + +template +std::map>, Tensor>> +LRI::constract_cvc_ds( + const std::map>>>>& cvc) +{ + const Data_Pack_Wrapper data_wrapper(this->data_pool, this->data_ab_name); + const LRI_Cal_Tools tools(this->period, this->data_pool, this->data_ab_name); + + const std::vector list_I = LRI_Cal_Aux::filter_list_map(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).a01, data_wrapper(Label::ab::a).Ds_ab); + const std::vector list_J = LRI_Cal_Aux::filter_list_map(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).b01, data_wrapper(Label::ab::b).Ds_ab); + const std::vector list_K0 = LRI_Cal_Aux::filter_list_set(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).a2, data_wrapper(Label::ab::a).index_Ds_ab[0]); + const std::vector list_L0 = LRI_Cal_Aux::filter_list_set(this->parallel->list_A.at(Label::Aab_Aab::a01b01_a2b2).b2, data_wrapper(Label::ab::b).index_Ds_ab[0]); + + const std::vector list_K = LRI_Cal_Aux::filter_list_map(list_K0, data_wrapper(Label::ab::a2b2).Ds_ab); + const std::vector list_L = LRI_Cal_Aux::filter_list_set(list_L0, data_wrapper(Label::ab::a2b2).index_Ds_ab[0]); + + std::map>> Hs; + for (auto& map_i : cvc) + { + const TA& I = map_i.first; + if (this->filter_atom->filter_for1(Label::ab_ab::a0b0_a2b2, I)) continue; // restrict I in the irreducible sector + for (auto& map_j : map_i.second) + { + const TAC& J = map_j.first; + if (this->filter_atom->filter_for32(Label::ab_ab::a0b0_a2b2, I, J, J)) continue; // restrict (I, J) in the irreducible sector + // init H_IJ + for (auto& map_k : map_j.second) + { + const TA& K = map_k.first; + for (auto& map_l : map_k.second) + { + const TAC& L = map_l.first; + const Tensor& CVC_IJKL = map_l.second; + const Tensor& D_KL = tools.get_Ds_ab(Label::ab::a2b2, K, L); + if (D_KL.empty()) continue; + // Hs[I][J] += Tensor_Multiply::gemv(CVC_IJKL, D_KL) + LRI_Cal_Aux::add_Ds(Tensor_Multiply::gemv(CVC_IJKL, D_KL), Hs[I][J]); + } + } + } + } + return Hs; +} + +} // end namespace RI + diff --git a/include/RI/ri/LRI-cal_cvc_mo.hpp b/include/RI/ri/LRI-cal_cvc_mo.hpp new file mode 100644 index 0000000..62dda1c --- /dev/null +++ b/include/RI/ri/LRI-cal_cvc_mo.hpp @@ -0,0 +1,496 @@ +#pragma once + +#include +#include "LRI.h" +#include "LRI_Cal_Aux.h" +#include "../global/Array_Operator.h" +#include "../global/Tensor_Multiply.h" +#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::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; + + // 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=" << GlobalV::MY_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=" << GlobalV::MY_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] = std::move(CI_JR.second * Global_Func::convert(phase)); + } + else { Ck_I_thread[iat2] += CI_JR.second * 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::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::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-cal_hartree.hpp b/include/RI/ri/LRI-cal_hartree.hpp new file mode 100644 index 0000000..57e375d --- /dev/null +++ b/include/RI/ri/LRI-cal_hartree.hpp @@ -0,0 +1,232 @@ +#pragma once + +#include +#include "LRI.h" +#include "LRI_Cal_Aux.h" +#include "../global/Array_Operator.h" +#include "../global/Tensor_Multiply.h" +#include +#include +#ifdef __MKL_RI +#include +#endif + +namespace RI +{ +template +std::map>>> // H(s,t)[k] +LRI::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) +{ +#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("Vs_").Ds_ab; + const std::map>>& Cs = this->data_pool.at("Cs_").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("Vs_"); + LRI_Cal_Aux::destroy_lock_result(lock_csk_result_add_map, Csk); + this->free_tensors_map2("Cs_"); + std::cout << "finished FT VsR and CsR to Vq and Csk." << std::endl; + } + #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_map(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_map(Csk, u, v); + const std::map>& Csk_v_u = Global_Func::find_map(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_map(Csk, s, t); + const std::map>& Csk_t_s = Global_Func::find_map(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 \ No newline at end of file diff --git a/include/RI/ri/LRI.h b/include/RI/ri/LRI.h index 5cf0c1b..4dafbd4 100644 --- a/include/RI/ri/LRI.h +++ b/include/RI/ri/LRI.h @@ -34,6 +34,7 @@ class LRI using TAC = std::pair; using Tdata_real = Global_Func::To_Real_t; using Tatom_pos = std::array; // tmp + using Tk = std::array; LRI(); @@ -64,6 +65,53 @@ class LRI std::map>> &Ds_result, const double fac_add_Ds = 1.0); + std::map>>>> cal_cvc(); + + std::map>> constract_cvc_ds( + const std::map>>>>& cvc); + + 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); + public: std::shared_ptr> parallel = std::make_shared>(); @@ -94,3 +142,6 @@ class LRI #include "LRI.hpp" #include "LRI-set.hpp" #include "LRI-cal_loop3.hpp" +#include "LRI-cal_cvc.hpp" +#include "LRI-cal_cvc_mo.hpp" +#include "LRI-cal_hartree.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..7508320 100644 --- a/include/RI/ri/LRI_Cal_Aux.h +++ b/include/RI/ri/LRI_Cal_Aux.h @@ -56,7 +56,14 @@ namespace LRI_Cal_Aux else D_result = D_result + D_add; } - + template + void add_Ds(const Tensor& D_in, Tensor& D_out, const Tdata fac) + { + if (D_out.empty()) + D_out = fac * D_in; + else + D_out += fac * D_in; + } template inline void add_Ds( Tensor &&D_add, @@ -160,16 +167,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 +192,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 +349,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 +386,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..9b01539 100644 --- a/include/RI/ri/LRI_Cal_Tools.h +++ b/include/RI/ri/LRI_Cal_Tools.h @@ -58,8 +58,16 @@ 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 const Tensor& get_Ds_ab( + const Label::ab& label, + const TAC& Aa, const TA& Ab) const + { + using namespace Array_Operator; + return Global_Func::find( + *this->Ds_ab_ptr.at(label), + Aa.first, TAC{ Ab, (-Aa.second) % this->period }); + } + inline std::vector split_b01(const Label::ab_ab &label) const { switch(label) { diff --git a/unittests/Test_All.hpp b/unittests/Test_All.hpp index e34f3e1..446eaf9 100644 --- a/unittests/Test_All.hpp +++ b/unittests/Test_All.hpp @@ -21,6 +21,7 @@ #include "unittests/physics/Exx-test.hpp" #include "unittests/physics/RPA-test.hpp" #include "unittests/physics/GW-test.hpp" +#include "unittests/physics/LR-test.hpp" namespace Test_All { @@ -92,5 +93,7 @@ namespace Test_All GW_Test::main(argc, argv); GW_Test::main>(argc, argv); GW_Test::main>(argc, argv); + + LR_Test::main(argc, argv); } } \ No newline at end of file diff --git a/unittests/physics/LR-ref b/unittests/physics/LR-ref new file mode 100644 index 0000000..7ed7a77 --- /dev/null +++ b/unittests/physics/LR-ref @@ -0,0 +1,59 @@ +3 +0 1 +0 3 2 3 3 6.000000000e-01 1.125000000e-01 1.200000000e+00 3.375000000e-01 1.800000000e+00 5.625000000e-01 6.000000000e-01 1.125000000e-01 1.200000000e+00 3.375000000e-01 1.800000000e+00 5.625000000e-01 6.000000000e-01 1.125000000e-01 1.200000000e+00 3.375000000e-01 1.800000000e+00 5.625000000e-01 6.000000000e-01 1.125000000e-01 1.200000000e+00 3.375000000e-01 1.800000000e+00 5.625000000e-01 6.000000000e-01 1.125000000e-01 1.200000000e+00 3.375000000e-01 1.800000000e+00 5.625000000e-01 6.000000000e-01 1.125000000e-01 1.200000000e+00 3.375000000e-01 1.800000000e+00 5.625000000e-01 +1 1 +0 3 2 3 3 1.200000000e+00 3.375000000e-01 2.400000000e+00 1.012500000e+00 3.600000000e+00 1.687500000e+00 1.200000000e+00 3.375000000e-01 2.400000000e+00 1.012500000e+00 3.600000000e+00 1.687500000e+00 1.200000000e+00 3.375000000e-01 2.400000000e+00 1.012500000e+00 3.600000000e+00 1.687500000e+00 1.200000000e+00 3.375000000e-01 2.400000000e+00 1.012500000e+00 3.600000000e+00 1.687500000e+00 1.200000000e+00 3.375000000e-01 2.400000000e+00 1.012500000e+00 3.600000000e+00 1.687500000e+00 1.200000000e+00 3.375000000e-01 2.400000000e+00 1.012500000e+00 3.600000000e+00 1.687500000e+00 +2 1 +0 3 2 3 3 1.800000000e+00 5.625000000e-01 3.600000000e+00 1.687500000e+00 5.400000000e+00 2.812500000e+00 1.800000000e+00 5.625000000e-01 3.600000000e+00 1.687500000e+00 5.400000000e+00 2.812500000e+00 1.800000000e+00 5.625000000e-01 3.600000000e+00 1.687500000e+00 5.400000000e+00 2.812500000e+00 1.800000000e+00 5.625000000e-01 3.600000000e+00 1.687500000e+00 5.400000000e+00 2.812500000e+00 1.800000000e+00 5.625000000e-01 3.600000000e+00 1.687500000e+00 5.400000000e+00 2.812500000e+00 1.800000000e+00 5.625000000e-01 3.600000000e+00 1.687500000e+00 5.400000000e+00 2.812500000e+00 +3 +0 3 +0 4 1 2 1 2 6.948874512e-02 5.135575125e-19 1.050820752e-01 -3.018515625e-03 1.050820752e-01 3.018515625e-03 1.590380420e-01 7.362912363e-19 +1 4 1 2 1 2 2.963159561e-01 1.074135938e-02 4.485606768e-01 3.371625000e-03 4.512944268e-01 3.406935938e-02 6.839351279e-01 3.191653125e-02 +2 4 1 2 1 2 6.839351279e-01 3.191653125e-02 1.035644880e+00 1.855532813e-02 1.044392880e+00 9.263995313e-02 1.583373050e+00 9.472443750e-02 +1 3 +0 4 1 2 1 2 2.963159561e-01 -1.074135938e-02 4.512944268e-01 -3.406935938e-02 4.485606768e-01 -3.371625000e-03 6.839351279e-01 -3.191653125e-02 +1 4 1 2 1 2 1.265219604e+00 -5.076299103e-17 1.929689341e+00 -7.551984375e-02 1.929689341e+00 7.551984375e-02 2.947633901e+00 -4.377345708e-18 +2 4 1 2 1 2 2.921389901e+00 3.037879688e-02 4.457462669e+00 -1.280420156e-01 4.467850919e+00 2.335989375e-01 6.828226198e+00 8.959865625e-02 +2 3 +0 4 1 2 1 2 6.839351279e-01 -3.191653125e-02 1.044392880e+00 -9.263995313e-02 1.035644880e+00 -1.855532813e-02 1.583373050e+00 -9.472443750e-02 +1 4 1 2 1 2 2.921389901e+00 -3.037879688e-02 4.467850919e+00 -2.335989375e-01 4.457462669e+00 1.280420156e-01 6.828226198e+00 -8.959865625e-02 +2 4 1 2 1 2 6.746213698e+00 2.231999308e-18 1.032186900e+01 -4.321033594e-01 1.032186900e+01 4.321033594e-01 1.582038425e+01 3.275044819e-16 +3 +0 3 +0 4 1 2 1 2 6.839524512e-02 -1.227909375e-02 1.028950752e-01 -2.153967188e-02 1.028950752e-01 -2.153967188e-02 1.546640420e-01 -3.704231250e-02 +1 4 1 2 1 2 2.897549561e-01 -6.293320313e-02 4.354386768e-01 -1.077553125e-01 4.381724268e-01 -1.132797656e-01 6.576911279e-01 -1.903373438e-01 +2 4 1 2 1 2 6.675326279e-01 -1.522698750e-01 1.002839880e+00 -2.592620156e-01 1.011587880e+00 -2.757328594e-01 1.517763050e+00 -4.609102500e-01 +1 3 +0 4 1 2 1 2 2.897549561e-01 -6.293320313e-02 4.381724268e-01 -1.132797656e-01 4.354386768e-01 -1.077553125e-01 6.576911279e-01 -1.903373438e-01 +1 4 1 2 1 2 1.225853604e+00 -3.131510625e-01 1.850957341e+00 -5.507822813e-01 1.850957341e+00 -5.507822813e-01 2.790169901e+00 -9.505248750e-01 +2 4 1 2 1 2 2.822974901e+00 -7.524988594e-01 4.260632669e+00 -1.316198109e+00 4.271020919e+00 -1.332156375e+00 6.434566198e+00 -2.286713531e+00 +2 3 +0 4 1 2 1 2 6.675326279e-01 -1.522698750e-01 1.011587880e+00 -2.757328594e-01 1.002839880e+00 -2.592620156e-01 1.517763050e+00 -4.609102500e-01 +1 4 1 2 1 2 2.822974901e+00 -7.524988594e-01 4.271020919e+00 -1.332156375e+00 4.260632669e+00 -1.316198109e+00 6.434566198e+00 -2.286713531e+00 +2 4 1 2 1 2 6.500176198e+00 -1.805300156e+00 9.829793997e+00 -3.178496953e+00 9.829793997e+00 -3.178496953e+00 1.483623425e+01 -5.492787188e+00 +3 +0 3 +0 4 1 2 1 2 6.948874512e-02 4.043396477e-19 1.050820752e-01 -3.018515625e-03 1.050820752e-01 3.018515625e-03 1.590380420e-01 9.254072118e-19 +1 4 1 2 1 2 2.963159561e-01 1.074135938e-02 4.485606768e-01 3.371625000e-03 4.512944268e-01 3.406935938e-02 6.839351279e-01 3.191653125e-02 +2 4 1 2 1 2 6.839351279e-01 3.191653125e-02 1.035644880e+00 1.855532813e-02 1.044392880e+00 9.263995313e-02 1.583373050e+00 9.472443750e-02 +1 3 +0 4 1 2 1 2 2.963159561e-01 -1.074135938e-02 4.512944268e-01 -3.406935938e-02 4.485606768e-01 -3.371625000e-03 6.839351279e-01 -3.191653125e-02 +1 4 1 2 1 2 1.265219604e+00 -7.062221901e-18 1.929689341e+00 -7.551984375e-02 1.929689341e+00 7.551984375e-02 2.947633901e+00 -1.645314744e-17 +2 4 1 2 1 2 2.921389901e+00 3.037879687e-02 4.457462669e+00 -1.280420156e-01 4.467850919e+00 2.335989375e-01 6.828226198e+00 8.959865625e-02 +2 3 +0 4 1 2 1 2 6.839351279e-01 -3.191653125e-02 1.044392880e+00 -9.263995313e-02 1.035644880e+00 -1.855532813e-02 1.583373050e+00 -9.472443750e-02 +1 4 1 2 1 2 2.921389901e+00 -3.037879688e-02 4.467850919e+00 -2.335989375e-01 4.457462669e+00 1.280420156e-01 6.828226198e+00 -8.959865625e-02 +2 4 1 2 1 2 6.746213698e+00 -1.974377693e-17 1.032186900e+01 -4.321033594e-01 1.032186900e+01 4.321033594e-01 1.582038425e+01 -4.630065862e-17 +3 +0 3 +0 4 1 2 1 2 6.839524512e-02 -1.227909375e-02 1.028950752e-01 -2.153967188e-02 1.028950752e-01 -2.153967188e-02 1.546640420e-01 -3.704231250e-02 +1 4 1 2 1 2 2.897549561e-01 -6.293320313e-02 4.354386768e-01 -1.077553125e-01 4.381724268e-01 -1.132797656e-01 6.576911279e-01 -1.903373438e-01 +2 4 1 2 1 2 6.675326279e-01 -1.522698750e-01 1.002839880e+00 -2.592620156e-01 1.011587880e+00 -2.757328594e-01 1.517763050e+00 -4.609102500e-01 +1 3 +0 4 1 2 1 2 2.897549561e-01 -6.293320313e-02 4.381724268e-01 -1.132797656e-01 4.354386768e-01 -1.077553125e-01 6.576911279e-01 -1.903373438e-01 +1 4 1 2 1 2 1.225853604e+00 -3.131510625e-01 1.850957341e+00 -5.507822813e-01 1.850957341e+00 -5.507822813e-01 2.790169901e+00 -9.505248750e-01 +2 4 1 2 1 2 2.822974901e+00 -7.524988594e-01 4.260632669e+00 -1.316198109e+00 4.271020919e+00 -1.332156375e+00 6.434566198e+00 -2.286713531e+00 +2 3 +0 4 1 2 1 2 6.675326279e-01 -1.522698750e-01 1.011587880e+00 -2.757328594e-01 1.002839880e+00 -2.592620156e-01 1.517763050e+00 -4.609102500e-01 +1 4 1 2 1 2 2.822974901e+00 -7.524988594e-01 4.271020919e+00 -1.332156375e+00 4.260632669e+00 -1.316198109e+00 6.434566198e+00 -2.286713531e+00 +2 4 1 2 1 2 6.500176198e+00 -1.805300156e+00 9.829793997e+00 -3.178496953e+00 9.829793997e+00 -3.178496953e+00 1.483623425e+01 -5.492787188e+00 diff --git a/unittests/physics/LR-test.hpp b/unittests/physics/LR-test.hpp new file mode 100644 index 0000000..ad9de68 --- /dev/null +++ b/unittests/physics/LR-test.hpp @@ -0,0 +1,218 @@ +#pragma once + +#define LR_REF_PRECISION 9 // significant digits when writing reference file +#define LR_REF_TOLERANCE 1e-7 // tolerance when comparing against reference + +#include "RI/physics/LR.h" +#include +#include +#include +#include +#include +#include + +namespace LR_Test +{ + +using Tdata = std::complex; +using TA = int; +using Tcell = int; +constexpr int Ndim = 3; +using TAC = std::pair>; +using Tk = std::array; + +static const std::string ref_path = "unittests/physics/LR-ref"; + +// -------- helpers: serialize / deserialize 2-level maps of arbitrary-rank tensors -------- +using TensorMap = std::map>>;// map> or map> + +static void write_tensor_map(std::ofstream &ofs, const TensorMap &m) +{ + ofs << m.size() << "\n"; + for (auto &[k1, inner] : m) { + ofs << k1 << " " << inner.size() << "\n"; + for (auto &[k2, t] : inner) { + ofs << k2 << " " << t.shape.size(); + for (auto d : t.shape) ofs << " " << d; + std::size_t total = 1; + for (auto d : t.shape) total *= d; + for (std::size_t i = 0; i < total; ++i) + ofs << " " << std::setprecision(LR_REF_PRECISION) << std::scientific + << t.ptr()[i].real() << " " << t.ptr()[i].imag(); + ofs << "\n"; + } + } +} + +static void read_tensor_map(std::ifstream &ifs, TensorMap &m) +{ + m.clear(); + std::size_t n1; ifs >> n1; + for (std::size_t i = 0; i < n1; ++i) { + int k1; std::size_t n2; + ifs >> k1 >> n2; + for (std::size_t j = 0; j < n2; ++j) { + int k2; std::size_t ndim; + ifs >> k2 >> ndim; + std::vector shape(ndim); + std::size_t total = 1; + for (std::size_t d = 0; d < ndim; ++d) { + ifs >> shape[d]; total *= shape[d]; + } + RI::Tensor t(shape); + for (std::size_t p = 0; p < total; ++p) { + double re, im; ifs >> re >> im; t.ptr()[p] = Tdata(re, im); + } + m[k1][k2] = std::move(t); + } + } +} + +static void compare_tensor_map(const TensorMap &a, const TensorMap &b) +{ + assert(a.size() == b.size()); + for (auto &[k1, inner_a] : a) { + assert(b.count(k1)); + auto &inner_b = b.at(k1); + assert(inner_a.size() == inner_b.size()); + for (auto &[k2, t_a] : inner_a) { + assert(inner_b.count(k2)); + auto &t_b = inner_b.at(k2); + assert(t_a.shape.size() == t_b.shape.size()); + std::size_t total = 1; + for (std::size_t d = 0; d < t_a.shape.size(); ++d) { + assert(t_a.shape[d] == t_b.shape[d]); + total *= t_a.shape[d]; + } + for (std::size_t p = 0; p < total; ++p) + assert(std::abs(t_a.ptr()[p] - t_b.ptr()[p]) < LR_REF_TOLERANCE); + } + } +} + +// -------- 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}}; + + // k1 = k2 = k_indices = all 3 + std::vector k_indices(nk), k1_indices(nk), k2_indices(nk); + for (int i = 0; i < nk; ++i) + k_indices[i] = k1_indices[i] = k2_indices[i] = i; + + // build q_list / q2kpair from k1-k2 differences + std::vector q_list; + std::map>> q2kpair; + { + auto mod1 = [](double x) { double r = x - std::floor(x); return r < 0 ? r+1.0 : r; }; + for (int k1 : k1_indices) + for (int k2 : 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])}}; + q2kpair[q].emplace_back(k1, k2); + } + for (auto& p : q2kpair) q_list.push_back(p.first); + } + + // Cs: shape {nabf=2, nw=3, nw=3}, all 1.0+0i + std::map>> Cs; + Cs[atom0][{atom0, {{0,0,0}}}] = 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, {{0,0,0}}}](mu,s,t) = Tdata(1.0, 0.0); + + // Vs / Ws: shape {nabf=2, nabf=2}, diag=0.01, off-diag=0.005 + auto make_V = [&]() { + std::map>> V; + V[atom0][{atom0, {{0,0,0}}}] = RI::Tensor({nabf, nabf}); + for (int i = 0; i < nabf; ++i) + for (int j = 0; j < nabf; ++j) + V[atom0][{atom0, {{0,0,0}}}](i,j) = (i==j) ? Tdata(0.01, 0.0) : Tdata(0.005, 0.0); + return V; + }; + auto Vs = make_V(); + auto Ws = make_V(); + + // map_psi: for each k, shape {nband=3, nw=3}, complex values + std::map>> map_psi; + for (int ik : k_indices) { + map_psi[ik][atom0] = RI::Tensor({nband, nw}); + for (int m = 0; m < nband; ++m) + for (int t = 0; t < nw; ++t) + 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)); + } + + // ====== init LR ====== + RI::LR lr; + std::array,3> latvec = {{{{1,0,0}}, {{0,1,0}}, {{0,0,1}}}}; + lr.set_parallel(MPI_COMM_WORLD, {{{atom0, {{0,0,0}}}}}, latvec, {{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}); + + lr.k1_indices = k1_indices; + lr.k2_indices = k2_indices; + lr.k_indices = k_indices; + lr.list_I = {atom0}; lr.list_J = {atom0}; lr.list_IJ = {atom0}; + lr.kindex_map = kfrac; + lr.nocc = nocc; lr.nvirt = nvirt; + lr.map_psi = std::move(map_psi); + lr.q_list = q_list; lr.q2kpair = q2kpair; + + // 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); + + // ====== regression: write or compare reference file ====== + if (argc >= 2 && std::string(argv[1]) == "--write-ref") + { + std::ofstream ofs_ref(ref_path); + assert(ofs_ref && "Failed to open ref file for writing"); + ofs_ref << std::scientific << std::setprecision(LR_REF_PRECISION); + write_tensor_map(ofs_ref, lr.Csk_ao_mo); + write_tensor_map(ofs_ref, V_A); + write_tensor_map(ofs_ref, V_B); + write_tensor_map(ofs_ref, W_A); + write_tensor_map(ofs_ref, W_B); + std::cout << "LR_ref written to " << ref_path << std::endl; + } + else + { + std::ifstream ifs_ref(ref_path); + assert(ifs_ref && "LR_ref not found; run with --write-ref to generate"); + TensorMap ref; + read_tensor_map(ifs_ref, ref); compare_tensor_map(lr.Csk_ao_mo, ref); + read_tensor_map(ifs_ref, ref); compare_tensor_map(V_A, ref); + read_tensor_map(ifs_ref, ref); compare_tensor_map(V_B, ref); + read_tensor_map(ifs_ref, ref); compare_tensor_map(W_A, ref); + read_tensor_map(ifs_ref, ref); compare_tensor_map(W_B, ref); + } + + lr.free_Cs(); + lr.free_Vs(); + lr.free_Ws(); + + MPI_Finalize(); +} + +} \ No newline at end of file diff --git a/unittests/ri/Cell_Nearest-test.hpp b/unittests/ri/Cell_Nearest-test.hpp index b0f94c6..f1ee4c6 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 +} \ No newline at end of file