summing over components

src/Hamiltonians/correlation_functions.jl

@@ -123,7 +123,21 @@ struct G2RealSpace{A,B,G<:Geometry,S} <: AbstractHamiltonian{S}
     geometry::G
     init::S
 end
-function G2RealSpace(geometry::Geometry, source=1, target=source)
+function G2RealSpace(
+    geometry::Geometry, source::Int=1, target::Int=source; sum_components=false
+)
+    if source < 1 || target < 1
+        throw(ArgumentError("`source` and `target` must be positive integers"))
+    end
+    if sum_components
+        if source ≠ 1 || target ≠ 1
+            throw(
+                ArgumentError("`source` or `target` can't be set if `sum_components=true`")
+            )
+        end
+        source = target = 0
+    end
+
     init = zeros(SArray{Tuple{size(geometry)...}})
     return G2RealSpace{source,target,typeof(geometry),typeof(init)}(geometry, init)
 end
@@ -136,37 +150,38 @@ LOStructure(::Type{<:G2RealSpace}) = IsDiagonal()
 
 num_offdiagonals(g2::G2RealSpace, _) = 0
 
-@inline function diagonal_element(
-    g2::G2RealSpace{A,B}, addr1::SingleComponentFockAddress, addr2::SingleComponentFockAddress
+@inline function _g2_diagonal_element(
+    g2::G2RealSpace{A,B}, onr1::SArray, onr2::SArray
 ) where {A, B}
     geo = g2.geometry
     result = g2.init
-    if addr1 ≡ addr2
-        v1 = v2 = onr(addr1, geo)
-    else
-        v1 = onr(addr1, geo)
-        v2 = onr(addr2, geo)
-    end
 
     @inbounds for i in eachindex(result)
         res_i = 0.0
         δ_vec = Offsets(geo)[i]
 
         # Case of n_i(n_i - 1) on the same component
         if A == B && all(==(0), δ_vec)
-            v1_offset = max.(v1 .- 1, 0)
-            result = setindex(result, dot(v2, v1_offset), i)
+            onr1_offset = max.(onr1 .- 1, 0)
+            result = setindex(result, dot(onr2, onr1_offset), i)
         else
-            result = setindex(result, circshift_dot(v2, v1, δ_vec), i)
+            result = setindex(result, circshift_dot(onr2, onr1, δ_vec), i)
         end
     end
     return result ./ length(geo)
 end
 function diagonal_element(g2::G2RealSpace{A,A}, addr::SingleComponentFockAddress) where {A}
-    return diagonal_element(g2, addr, addr)
+    onr1 = onr(addr, g2.geometry)
+    return _g2_diagonal_element(g2, onr1, onr1)
 end
 function diagonal_element(g2::G2RealSpace{A,B}, addr::CompositeFS) where {A,B}
-    return diagonal_element(g2, addr.components[A], addr.components[B])
+    onr1 = onr(addr.components[A], g2.geometry)
+    onr2 = onr(addr.components[B], g2.geometry)
+    return _g2_diagonal_element(g2, onr1, onr2)
+end
+function diagonal_element(g2::G2RealSpace{0,0}, addr::CompositeFS)
+    onr1 = sum(onr, addr.components)
+    return _g2_diagonal_element(g2, onr1, onr1)
 end
 
 """

test/Hamiltonians.jl

@@ -166,7 +166,6 @@ using Rimu.Hamiltonians: UnitVectors, Offsets
             @test length(geom) == prod(dims)
             @test Rimu.Hamiltonians.fold(geom) == fold
             @test eval(Meta.parse(repr(geom))) == geom
-            @test num_dimensions(geom) == length(dims)
         end
     end
 
@@ -867,13 +866,29 @@ using Rimu.Hamiltonians: circshift_dot
 
     @testset "G2RealSpace" begin
         @testset "1D G2RealCorrelator comparison" begin
-            addr = near_uniform(BoseFS{6,6})
-            H = HubbardReal1D(addr)
-            v = normalize!(H * (H * (H * DVec(addr => 1.0))))
+            @testset "single components" begin
+                addr = near_uniform(BoseFS{6,6})
+                H = HubbardReal1D(addr)
+                v = normalize!(H * (H * (H * DVec(addr => 1.0))))
+
+                g2 = dot(v, G2RealSpace(Geometry(6)), v)
+                for d in 0:5
+                    @test g2[d + 1] ≈ dot(v, G2RealCorrelator(d), v)
+                end
+            end
+
+            @testset "sum of components" begin
+                addr = CompositeFS(BoseFS(4, 1=>1), BoseFS(4, 1=>1), BoseFS(4, 1=>1))
+                H = HubbardRealSpace(addr)
+                v = normalize!(H * (H * (H * DVec(addr => 1.0))))
+
+                g2 = dot(v, G2RealSpace(Geometry(4); sum_components=true), v)
+                for d in 0:3
+                    @test g2[d + 1] ≈ dot(v, G2RealCorrelator(d), v)
+                end
 
-            g2 = dot(v, G2RealSpace(Geometry(6)), v)
-            for d in 0:5
-                @test g2[d + 1] ≈ dot(v, G2RealCorrelator(d), v)
+                g2_nosum = dot(v, G2RealSpace(Geometry(4)), v)
+                @test iszero(g2_nosum)
             end
         end
 
@@ -888,11 +903,18 @@ using Rimu.Hamiltonians: circshift_dot
             g2_21 = G2RealSpace(geom, 2, 1)
             g2_22 = G2RealSpace(geom, 2, 2)
 
+            # Sum is N1 * (N2 - δ_12) / M
             @test sum(diagonal_element(g2_11, addr)) == 0
             @test sum(diagonal_element(g2_12, addr)) == 3.5
             @test sum(diagonal_element(g2_21, addr)) == 3.5
             @test sum(diagonal_element(g2_22, addr)) == 70
+
+            # Swapping components flips axes
+            @test diagonal_element(g2_11, addr) == [0 0; 0 0; 0 0]
             @test diagonal_element(g2_12, addr) == [1 4; 2 5; 3 6] ./ 6
+            @test diagonal_element(g2_21, addr) == [1 4; 3 6; 2 5] ./ 6
+
+
         end
     end