Refactor: itiaiw2iat interface in UnitCell (#2241)

* Refactor: itiaiw2iat interface in UnitCell

* Fix: compiling error in src_ri

---------

Co-authored-by: dyzheng <[email protected]>

Author: dyzheng

source/module_cell/test/unitcell_test_readpp.cpp

@@ -61,6 +61,12 @@ Magnetism::~Magnetism()
  *     - cal_nwfc(): calcuate the total number of local basis: NSPIN != 4
  *     - this corresponds to number_of_proj, PP_BETA in pp file, and atoms[it].l_nchi[nw], nw from orb file
  *     - setup GlobalV::NLOCAL
+ *     - interfaces initialed in this function:
+ *       - itia2iat
+ *       - iat2iwt
+ *       - itiaiw2iwt
+ *       - iwt2iat
+ *       - iwt2iw
  *   - CalNwfc2
  *     - cal_nwfc(): calcuate the total number of local basis: NSPIN == 4
  *   - CheckStructure
@@ -257,6 +263,37 @@ TEST_F(UcellTest,CalNwfc1)
 	EXPECT_EQ(ucell->atoms[1].nw,9);
 	EXPECT_EQ(ucell->nwmax,9);
 	EXPECT_EQ(GlobalV::NLOCAL,3*9);
+	//check itia2iat
+	EXPECT_EQ(ucell->itia2iat.getSize(), 4);
+	EXPECT_EQ(ucell->itia2iat(0,0), 0);
+	EXPECT_EQ(ucell->itia2iat(0,1), 0);
+	EXPECT_EQ(ucell->itia2iat(1,0), 1);
+	EXPECT_EQ(ucell->itia2iat(1,1), 2);
+	//check iat2iwt
+	EXPECT_EQ(ucell->iat2iwt.size(), 3);
+	EXPECT_EQ(ucell->iat2iwt[0], 0);
+	EXPECT_EQ(ucell->iat2iwt[1], 9);
+	EXPECT_EQ(ucell->iat2iwt[2], 18);
+	//check itiaiw2iwt
+	EXPECT_EQ(ucell->itiaiw2iwt(0, 0, 0), 0);
+	EXPECT_EQ(ucell->itiaiw2iwt(0, 0, 1), 1);
+	EXPECT_EQ(ucell->itiaiw2iwt(0, 0, 8), 8);
+	EXPECT_EQ(ucell->itiaiw2iwt(1, 0, 0), 9);
+	EXPECT_EQ(ucell->itiaiw2iwt(1, 1, 0), 18);
+	//check itia2iat
+	EXPECT_EQ(ucell->itia2iat.getSize(), 4);
+	EXPECT_EQ(ucell->itia2iat(0, 0), 0);
+	EXPECT_EQ(ucell->itia2iat(0, 1), 0);
+	EXPECT_EQ(ucell->itia2iat(1, 0), 1);
+	EXPECT_EQ(ucell->itia2iat(1, 1), 2);
+	//check iwt2iat
+	EXPECT_EQ(ucell->iwt2iat[0], 0);
+	EXPECT_EQ(ucell->iwt2iat[10], 1);
+	EXPECT_EQ(ucell->iwt2iat[20], 2);
+	//check iwt2iw
+	EXPECT_EQ(ucell->iwt2iw[0], 0);
+	EXPECT_EQ(ucell->iwt2iw[10], 1);
+	EXPECT_EQ(ucell->iwt2iw[20], 2);
 }
 
 TEST_F(UcellTest,CalNwfc2)

source/module_cell/unitcell.cpp

@@ -39,7 +39,6 @@ UnitCell::UnitCell()
 
     itia2iat.create(1, 1);
     lc = new int[3];
-    itiaiw2iwt.create(1, 1, 1);
 
     latvec = ModuleBase::Matrix3();
     latvec_supercell = ModuleBase::Matrix3();
@@ -873,7 +872,8 @@ void UnitCell::cal_nwfc(std::ofstream &log)
 	this->iwt2iw = new int[GlobalV::NLOCAL];
 
 	this->itia2iat.create(ntype, namax);
-	this->itiaiw2iwt.create(ntype, namax, nwmax*GlobalV::NPOL);
+	//this->itiaiw2iwt.create(ntype, namax, nwmax*GlobalV::NPOL);
+	this->iat2iwt.resize(nat);
 	int iat=0;
 	int iwt=0;
 	for(int it = 0;it < ntype;it++)
@@ -882,9 +882,10 @@ void UnitCell::cal_nwfc(std::ofstream &log)
 		{
 			this->itia2iat(it, ia) = iat;
 			//this->iat2ia[iat] = ia;
+			this->iat2iwt[iat] = iwt;
 			for(int iw=0; iw<atoms[it].nw * GlobalV::NPOL; iw++)
 			{
-				this->itiaiw2iwt(it, ia, iw) = iwt;
+				//this->itiaiw2iwt(it, ia, iw) = iwt;
 				this->iwt2iat[iwt] = iat;
 				this->iwt2iw[iwt] = iw;
 				++iwt;

source/module_cell/unitcell.h

@@ -67,7 +67,14 @@ class UnitCell
 	int *iwt2iat; // iwt ==> iat.
 	int *iwt2iw; // iwt ==> iw, Peize Lin add 2018-07-02
     ModuleBase::IntArray itia2iat;//(it, ia)==>iat, the index in nat, add 2009-3-2 by mohan
-    ModuleBase::IntArray itiaiw2iwt;//(it, ia, iw)==>iwt, the index in nwfc, add 2009-3-2 by mohan
+    //atom index iat to the first global index for orbital of this atom
+    std::vector<int> iat2iwt;
+    // indexing tool for find orbital global index from it,ia,iw
+    template<typename Tiait>
+    inline Tiait itiaiw2iwt(const Tiait &it, const Tiait &ia, const Tiait &iw) const
+    {
+        return Tiait(this->iat2iwt[this->itia2iat(it, ia)] + iw);
+    }
 
     //========================================================
     // indexing tools for ia and it

source/module_io/test/for_testing_input_conv.h

@@ -111,7 +111,6 @@ UnitCell::UnitCell(){
 
 	itia2iat.create(1, 1);
 	lc = new int[3];
-	itiaiw2iwt.create(1, 1, 1);
 
 	latvec = ModuleBase::Matrix3();
 	latvec_supercell = ModuleBase::Matrix3();

source/src_ri/exx_abfs-parallel-communicate-hexx-template.h

@@ -83,14 +83,14 @@ Tmatrix Exx_Abfs::Parallel::Communicate::Hexx::Ra2D_to_Km2D(const Parallel_Orbit
 	{
 		for(auto & HR_a2D_A : HR_a2D[is])
 		{
-			const size_t iat1 = HR_a2D_A.first;
-			const size_t it1 = GlobalC::ucell.iat2it[iat1];
-			const size_t ia1 = GlobalC::ucell.iat2ia[iat1];
+			const int iat1 = HR_a2D_A.first;
+			const int it1 = GlobalC::ucell.iat2it[iat1];
+			const int ia1 = GlobalC::ucell.iat2ia[iat1];
 			for(auto & HR_a2D_B : HR_a2D_A.second)
 			{
-				const size_t iat2 = HR_a2D_B.first;
-				const size_t it2 = GlobalC::ucell.iat2it[iat2];
-				const size_t ia2 = GlobalC::ucell.iat2ia[iat2];
+				const int iat2 = HR_a2D_B.first;
+				const int it2 = GlobalC::ucell.iat2it[iat2];
+				const int ia2 = GlobalC::ucell.iat2ia[iat2];
 
 				Tmatrix HK_a2D;
 				for(auto & HR_a2D_C : HR_a2D_B.second)

source/src_ri/exx_lcao.cpp

@@ -1026,7 +1026,7 @@ ofs_mpi.close();
 		if(GlobalV::GAMMA_ONLY_LOCAL)
 		{
 			std::ofstream ofs("LOC.DM.dat",std::ofstream::app);
-			const int it1=0, it2=0;
+			const size_t it1=0, it2=0;
 			for( size_t ia1=0; ia1!=GlobalC::ucell.atoms[it1].na; ++ia1 )
 				for( size_t ia2=0; ia2!=GlobalC::ucell.atoms[it2].na; ++ia2 )
 					for( size_t is=0; is!=GlobalV::NSPIN; ++is )