Fix accumulator aliasing and Hartree atom-list indexing

include/RI/ri/LRI_Cal_Aux.h

@@ -95,7 +95,7 @@ namespace LRI_Cal_Aux
 		if (D_result.empty())
 		{
 			if(Tdata(1.0)==fac)
-				D_result = D_add;
+				D_result = D_add.copy();
 			else if(Tdata(-1.0)==fac)
 				D_result = -D_add;
 			else

include/RI/ri/LRI_k-cal_cvc_mo.hpp

@@ -467,7 +467,7 @@ LRI_k<TA, Tcell, Ndim, Tdata>::cal_cvc_mo_k_hartree_onthefly(
 		{
 			for (int k1_index=0; k1_index<k1_indices.size(); ++k1_index)
 			{
-				const TA& mu = list_J[mu_index];
+				const TA& mu = list_I[mu_index];
 				const int k1 = k1_indices[k1_index];
 				// 2.1 calculate C^mu (i,a^*)[k1,k1] on-the-fly
 				const Tensor<Tdata> C_mu_ia = Cs_ao_mo_to_Cs_mo(Cs_ao_mo, map_psi, k1, k1,

include/RI/ri/LRI_k-cal_hartree.hpp

@@ -121,7 +121,7 @@ LRI_k<TA, Tcell, Ndim, Tdata>::cal_cvcd_k_hartree(
 #pragma omp for schedule(dynamic) collapse(2)
 		for (int v_index=0; v_index<list_I.size(); v_index++)
 		{
-			for (int u_index=0; u_index<list_I.size(); u_index++)
+			for (int u_index=0; u_index<list_J.size(); u_index++)
 			{
 				const TA v = list_I[v_index];
 				const TA u = list_J[u_index];

unittests/ri/LRI-k-atom-lists-test.hpp

@@ -0,0 +1,104 @@
+#pragma once
+
+#include "RI/ri/LRI_k.h"
+
+#include <cmath>
+#include <complex>
+#include <iostream>
+#include <stdexcept>
+
+namespace LRI_K_Atom_Lists_Test
+{
+using Scalar = std::complex<double>;
+using LRI = RI::LRI_k<int, int, 3, Scalar>;
+using Tensors = std::map<int, std::map<LRI::TAC, RI::Tensor<Scalar>>>;
+
+inline RI::Tensor<Scalar> tensor(const RI::Shape_Vector &shape, double value)
+{
+    RI::Tensor<Scalar> result(shape);
+    for (auto &element : *result.data)
+        element = value;
+    return result;
+}
+
+inline void setup(LRI &lri)
+{
+    lri.set_parallel(MPI_COMM_SELF, {{0, {0, 0, 0}}, {1, {0.3, 0, 0}}, {2, {0.7, 0, 0}}},
+                     {{{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}}, {1, 1, 1},
+                     RI::Global_Func::to_vector(RI::Label::array_ab_ab));
+    Tensors cs, vs;
+    for (int a = 0; a < 3; ++a)
+        for (int b = 0; b < 3; ++b)
+        {
+            cs[a][{b, {0, 0, 0}}] = tensor({1, 1, 1}, 1 + a + 2*b);
+            vs[a][{b, {0, 0, 0}}] = tensor({1, 1}, 2 + a + 3*b);
+        }
+    lri.set_tensors_map2(cs, {RI::Label::ab::a}, {}, "Cs");
+    lri.set_tensors_map2(vs, {RI::Label::ab::a0b0}, {}, "Vs");
+}
+
+inline void check(const RI::Tensor<Scalar> &actual, double expected)
+{
+    if (actual.empty() || actual.data->size() != 1)
+        throw std::runtime_error("Missing atom-list contraction output");
+    const double error = std::abs((*actual.data)[0] - expected);
+    if (!std::isfinite(error) || error > 1e-11 * (1 + std::abs(expected)))
+        throw std::runtime_error("Atom-list contraction differs from scalar reference");
+}
+
+inline void test_hartree()
+{
+    LRI lri;
+    setup(lri);
+    const std::vector<int> left = {0, 1, 2}, right = {0, 1};
+    std::map<int, std::map<int, std::map<int, RI::Tensor<Scalar>>>> density;
+    double m[3] = {}, n[3] = {};
+    for (int v : left)
+        for (int u : right)
+        {
+            const double d = 0.1 * (1 + 2*v + u);
+            density[v][u][0] = tensor({1, 1}, d);
+            m[u] += (1 + u + 2*v) * d;
+            m[v] += (1 + v + 2*u) * d;
+        }
+    for (int mu : left)
+        for (int nu : left)
+            n[mu] += (2 + mu + 3*nu) * m[nu];
+
+    const auto result = lri.cal_cvcd_k_hartree(
+        density, {{0, 0, 0}}, {0}, left, right, left, "Cs", "Vs");
+    for (int s : left)
+        for (int t : right)
+            check(result.at(s).at(t).at(0),
+                  (1 + s + 2*t) * n[s] + (1 + t + 2*s) * n[t]);
+    std::cout << "PASS Hartree with unequal left/right atom lists" << std::endl;
+}
+
+inline void test_cvc()
+{
+    LRI lri;
+    setup(lri);
+    std::map<int, std::map<int, RI::Tensor<Scalar>>> cs_mo, psi;
+    for (int atom = 0; atom < 3; ++atom)
+    {
+        cs_mo[0][atom] = tensor({1, 1, 2}, atom + 1);
+        psi[0][atom] = tensor({2, 1}, 1);
+    }
+    const std::vector<std::pair<std::vector<int>, std::vector<int>>> partitions = {
+        {{0}, {1}}, {{0}, {1, 2}}, {{0, 1}, {2}}};
+    for (const auto &partition : partitions)
+        for (bool is_A : {false, true})
+        {
+            double expected = 0;
+            for (int mu : partition.first)
+                for (int nu : partition.second)
+                    // Each MO coefficient contains its two conjugate contributions.
+                    expected += 4 * (mu + 1) * (2 + mu + 3*nu) * (nu + 1);
+            const auto result = lri.cal_cvc_mo_k_hartree_onthefly(
+                cs_mo, psi, {0}, {0}, partition.first, partition.second,
+                {"O", "V", "O", "V"}, 1, 1, "Vs", is_A);
+            check(result.at(0).at(0), expected);
+        }
+    std::cout << "PASS CVC with distinct and unequal left/right atom lists" << std::endl;
+}
+}

unittests/ri/LRI-loop3-test.hpp

@@ -13,6 +13,11 @@
 #include"unittests/print_stl.h"
 #include"unittests/global/Tensor-test.h"
 
+#include<cmath>
+#include<complex>
+#include<limits>
+#include<stdexcept>
+#include<vector>
 #include<map>
 #include<unordered_map>
 #include<iostream>
@@ -396,12 +401,124 @@ namespace LRI_Loop3_Test
 	}
 
 
+	template<typename Tdata>
+	void test_const_tensor_accumulation()
+	{
+		using Treal = RI::Global_Func::To_Real_t<Tdata>;
+		RI::Tensor<Tdata> input({2});
+		for(std::size_t i=0; i<2; ++i)
+			(*input.data)[i] = RI::Global_Func::convert<Tdata>(std::complex<Treal>(i+1, 0.5));
+		const RI::Tensor<Tdata> &source = input;
+		RI::Tensor<Tdata> first, second;
+		RI::LRI_Cal_Aux::add_Ds(source, first, Tdata(1));
+		RI::LRI_Cal_Aux::add_Ds(source, second, Tdata(1));
+		RI::LRI_Cal_Aux::add_Ds(source, first, Tdata(2));
+		for(std::size_t i=0; i<2; ++i)
+		{
+			const Tdata expected = RI::Global_Func::convert<Tdata>(std::complex<Treal>(i+1, 0.5));
+			if((*source.data)[i]!=expected || (*second.data)[i]!=expected
+				|| (*first.data)[i]!=Tdata(3)*expected)
+				throw std::runtime_error("Accumulating a const tensor changed its input or another result");
+		}
+		std::cout<<"PASS const tensor accumulation preserves input and independent results"<<std::endl;
+	}
+
+
+	template<typename Tdata>
+	void test_loop3_accumulation(
+		const RI::LRI_Cal_Mode cal_mode,
+		const std::vector<Tdata> &factors = {Tdata(0), Tdata(-1), Tdata(1), Tdata(1),
+			Tdata(-1), Tdata(0.5), Tdata(0), Tdata(-0.5), Tdata(0.3)})
+	{
+		using LRI = RI::LRI<int,int,1,Tdata>;
+		using Treal = typename LRI::Tdata_real;
+		using T_Ds = std::map<int, std::map<typename LRI::TAC, RI::Tensor<Tdata>>>;
+		const typename LRI::TAC key = {0,{0}};
+		const auto labels = RI::Global_Func::to_vector(RI::Label::array_ab_ab);
+		const auto make_tensor = [](const RI::Shape_Vector &shape)
+		{
+			RI::Tensor<Tdata> D(shape);
+			for(std::size_t i=0; i<D.data->size(); ++i)
+				(*D.data)[i] = RI::Global_Func::convert<Tdata>(
+					std::complex<Treal>(0.5+0.07*i, 0.1+0.03*i));
+			return D;
+		};
+
+		LRI lri;
+		lri.set_parallel(MPI_COMM_SELF, {{0,{0.0}}}, {{{1.0}}}, {1}, labels);
+		for(const RI::Label::ab label : RI::Label::array_ab)
+		{
+			const bool coefficient = label==RI::Label::ab::a || label==RI::Label::ab::b;
+			T_Ds Ds;
+			Ds[0][key] = make_tensor(coefficient ? RI::Shape_Vector{2,2,2} : RI::Shape_Vector{2,2});
+			lri.set_tensors_map2(Ds, {label});
+		}
+
+		const RI::Tensor<Tdata> seed = make_tensor({2,2});
+		const Treal tolerance = 128*std::numeric_limits<Treal>::epsilon();
+		std::vector<std::vector<RI::Label::ab_ab>> label_groups;
+		for(const RI::Label::ab_ab label : labels)
+			label_groups.push_back({label});
+		label_groups.push_back(labels);
+
+		for(const auto &cal_labels : label_groups)
+		{
+			T_Ds reference;
+			lri.cal_mode = RI::LRI_Cal_Mode::CPU;
+			lri.cal_loop3(cal_labels, reference);
+			const RI::Tensor<Tdata> contribution = reference.at(0).at(key).copy();
+			const Treal norm = contribution.norm(2);
+			if(!std::isfinite(norm) || norm==0)
+				throw std::runtime_error("Invalid CPU reference contraction");
+
+			lri.cal_mode = cal_mode;
+			for(const bool seeded : {false, true})
+			{
+				T_Ds result;
+				// Tensor copies share storage unless copy() is used explicitly.
+				if(seeded)
+					result[0][key] = seed.copy();
+				Tdata accumulated_factor = 0;
+				Treal accumulated_magnitude = 0;
+				for(const Tdata factor : factors)
+				{
+					lri.cal_loop3(cal_labels, result, factor);
+					accumulated_factor += factor;
+					accumulated_magnitude += std::abs(factor);
+					const RI::Tensor<Tdata> &actual = result.at(0).at(key);
+					if(result.size()!=1 || result.at(0).size()!=1 || actual.shape.size()!=2
+						|| actual.shape[0]!=2 || actual.shape[1]!=2)
+						throw std::runtime_error("Unexpected contraction output shape");
+					for(std::size_t i=0; i<seed.data->size(); ++i)
+					{
+						const Tdata initial = seeded ? (*seed.data)[i] : Tdata(0);
+						const Tdata expected = initial + accumulated_factor*(*contribution.data)[i];
+						const Treal error = std::abs((*actual.data)[i] - expected);
+						const Treal scale = 1 + std::abs(initial)
+							+ accumulated_magnitude*std::abs((*contribution.data)[i]);
+						if(!std::isfinite(error) || error>tolerance*scale)
+						{
+							std::cerr<<RI::Label_Tools::get_name(cal_labels)
+								<<" mode="<<static_cast<int>(cal_mode)<<" seeded="<<seeded
+								<<" factor="<<factor<<" element="<<i<<" error="<<error<<std::endl;
+							throw std::runtime_error("Accumulation differs from the CPU reference");
+						}
+					}
+				}
+			}
+		}
+		std::cout<<"Accumulation: all 18 contractions and their sum passed, mode="<<static_cast<int>(cal_mode)<<std::endl;
+	}
+
+
 	template<typename Tdata>
 	void main(int argc, char *argv[])
 	{
 		int mpi_init_provide;
 		MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &mpi_init_provide);
 
+		test_const_tensor_accumulation<Tdata>();
+
 		const std::vector<RI::LRI_Cal_Mode> cal_modes = {
 			RI::LRI_Cal_Mode::CPU,
 			RI::LRI_Cal_Mode::CPU_fine_grained_lock,
@@ -410,7 +527,16 @@ namespace LRI_Loop3_Test
 			#endif
 		};
 		for(const RI::LRI_Cal_Mode cal_mode: cal_modes)
+		{
 			test_loop3<Tdata>(cal_mode);
+			test_loop3_accumulation<Tdata>(cal_mode);
+			// Exercise imaginary coefficients, including the first write to an empty result.
+			using Treal = RI::Global_Func::To_Real_t<Tdata>;
+			test_loop3_accumulation<Tdata>(cal_mode, {
+				RI::Global_Func::convert<Tdata>(std::complex<Treal>(0.25, 0.5)),
+				RI::Global_Func::convert<Tdata>(std::complex<Treal>(0, -1)),
+				RI::Global_Func::convert<Tdata>(std::complex<Treal>(-0.25, 0.5))});
+		}
 
 		MPI_Finalize();
 	}