Feature/reduced density matrix (#289)

## New features
- The `n`-particle reduced density matrix operator is added to Rimu,
i.e., `ReducedDensityMatrix(n)`.
- New reduction function `sum_mutating!` allows parallel reduction of
array-valued data from `PDVec`s minimizing allocations.

## Modified behaviour
- Modified `Interfaces.dot_from_right(target, op, source::AbstractDVec)`
in `src/DictVectors/abstractdvec.jl`, which calculates the same thing
using `sum` operation.
- Removed `Interfaces.dot_from_right(target, op, source::PDVec)`
function, which was unnecessary since the same code already exists with
type-specific argument `source::AbstractDVec` in abstractdvec.jl script.

---------

Co-authored-by: Joachim Brand <[email protected]>
Co-authored-by: mtsch <[email protected]>

docs/src/dictvectors.md

@@ -23,6 +23,8 @@ localpart
 apply_operator!
 sort_into_targets!
 working_memory
+mapreduce
+sum_mutating!
 ```
 
 ## Supported operations

docs/src/hamiltonians.md

@@ -92,6 +92,7 @@ StringCorrelator
 DensityMatrixDiagonal
 SingleParticleExcitation
 TwoParticleExcitation
+ReducedDensityMatrix
 Momentum
 AxialAngularMomentumHO
 ```

src/DictVectors/DictVectors.jl

@@ -22,11 +22,11 @@ using ..Interfaces: Interfaces, AbstractDVec, AdjointUnknown,
     CompressionStrategy, IsDiagonal, LOStructure,
     apply_column!, apply_operator!, compress!,
     diagonal_element, offdiagonals, step_stats, AbstractHamiltonian, AbstractOperator,
-    dot_from_right
+    AbstractObservable, dot_from_right
 using ..StochasticStyles: StochasticStyles, IsDeterministic
 
 import ..Interfaces: deposit!, storage, StochasticStyle, default_style, freeze, localpart,
-    working_memory
+    working_memory, sum_mutating!
 
 export deposit!, storage, walkernumber, walkernumber_and_length, dot_from_right
 export DVec, InitiatorDVec, PDVec, PDWorkingMemory

src/DictVectors/abstractdvec.jl

@@ -263,7 +263,7 @@ walkernumber_and_length(v) = walkernumber(v), length(v)
 ###
 ### Vector-operator functions
 ###
-function LinearAlgebra.mul!(w::AbstractDVec, h::AbstractOperator, v::AbstractDVec)
+function LinearAlgebra.mul!(w::AbstractDVec, h::AbstractObservable, v::AbstractDVec)
     empty!(w)
     for (key, val) in pairs(v)
         w[key] += diagonal_element(h, key) * val
@@ -274,7 +274,7 @@ function LinearAlgebra.mul!(w::AbstractDVec, h::AbstractOperator, v::AbstractDVe
     return w
 end
 
-function Base.:*(h::AbstractOperator, v::AbstractDVec)
+function Base.:*(h::AbstractObservable, v::AbstractDVec)
     T = promote_type(scalartype(h), scalartype(v))
     if eltype(h) ≠ scalartype(h)
         throw(ArgumentError("Operators with non-scalar eltype don't support "*
@@ -286,13 +286,13 @@ function Base.:*(h::AbstractOperator, v::AbstractDVec)
 end
 
 # docstring in Interfaces
-function LinearAlgebra.dot(w::AbstractDVec, op::AbstractOperator, v::AbstractDVec)
+function LinearAlgebra.dot(w::AbstractDVec, op::AbstractObservable, v::AbstractDVec)
     return dot(LOStructure(op), w, op, v)
 end
-function LinearAlgebra.dot(::AdjointUnknown, w, op::AbstractOperator, v)
+function LinearAlgebra.dot(::AdjointUnknown, w, op::AbstractObservable, v)
     return dot_from_right(w, op, v)
 end
-function LinearAlgebra.dot(::LOStructure, w, op::AbstractOperator, v)
+function LinearAlgebra.dot(::LOStructure, w, op::AbstractObservable, v)
     if length(w) < length(v)
         return conj(dot_from_right(v, op', w)) # turn args around to execute faster
     else
@@ -301,14 +301,15 @@ function LinearAlgebra.dot(::LOStructure, w, op::AbstractOperator, v)
 end
 
 # docstring in Interfaces
-function Interfaces.dot_from_right(w, op, v::AbstractDVec)
-    T = typeof(zero(valtype(w)) * zero(eltype(op)) * zero(valtype(v)))
-    result = zero(T)
-    for (key, val) in pairs(v)
-        result += conj(w[key]) * diagonal_element(op, key) * val
-        for (add, elem) in offdiagonals(op, key)
-            result += conj(w[add]) * elem * val
+function Interfaces.dot_from_right(target, op, source::AbstractDVec)
+    T = typeof(zero(valtype(target)) * zero(valtype(source)) * zero(eltype(op)))
+
+    result = sum(pairs(source); init=zero(T)) do (k, v)
+        res = conj(target[k]) * diagonal_element(op, k) * v
+        for (k_off, v_off) in offdiagonals(op, k)
+            res += conj(target[k_off]) * v_off * v
         end
+        res
     end
-    return result
+    return result::T
 end

src/DictVectors/pdvec.jl

@@ -513,9 +513,9 @@ function Base.iterate(t::PDVecIterator, (segment_id, state))
 end
 
 """
-    mapreduce(f, op, keys(::PDVec); init)
-    mapreduce(f, op, values(::PDVec); init)
-    mapreduce(f, op, pairs(::PDVec); init)
+    mapreduce(f, op, keys(::PDVec); [init])
+    mapreduce(f, op, values(::PDVec); [init])
+    mapreduce(f, op, pairs(::PDVec); [init])
 
 Perform a parallel reduction operation on [`PDVec`](@ref)s. MPI-compatible. Is used in the
 definition of various functions from Base such as `reduce`, `sum`, `prod`, etc.
@@ -537,6 +537,38 @@ end
 # reduction operator.
 Base.sum(f, t::PDVecIterator; kwargs...) = mapreduce(f, +, t; kwargs...)
 
+"""
+    sum_mutating!(accumulator, [f! = add!], keys(::PDVec); [init])
+    sum_mutating!(accumulator, [f! = add!], values(::PDVec); [init])
+    sum_mutating!(accumulator, [f! = add!], pairs(::PDVec); [init])
+
+Perform a parallel sum on [`PDVec`](@ref)s for vector-valued results while minimizing
+allocations. The result of the sum will be added to `accumulator` and stored in
+`accumulator`. MPI-compatible. If `f!` is provided, it must accept two arguments, the first
+being the accumulator and the second the element of the iterator. Otherwise,`add!` is used.
+
+If provided, `init` must be a neutral element for `+` and of the same type as `accumulator`.
+
+See also [`mapreduce`](@ref).
+"""
+function Interfaces.sum_mutating!(accu, f!, iterator::PDVecIterator; kwargs...)
+    interim_result = _sum_mutating(accu, f!, iterator; kwargs...)
+    add!(accu, interim_result)
+    return accu
+end
+# I'm not sure why this function barrier is necessary, but it saves 10% cpu time (or 60ms)
+# in a benchmark with a PDVec of length 12870.
+function _sum_mutating(accu, f!, iterator::PDVecIterator; kwargs...)
+    interim_result = Folds.mapreduce( # parallel reduction via threads
+        +, Iterators.filter(!isempty, iterator.vector.segments); kwargs...
+    ) do segment
+        # each segment gets is own copy of the accumulator such that each thread can work
+        # on it independently. Later the accumulators are merged.
+        sum_mutating!(zero(accu), f!, iterator.selector(segment))
+    end
+    return merge_remote_reductions(iterator.vector.communicator, +, interim_result) # MPI
+end
+
 """
     all(predicate, keys(::PDVec); kwargs...)
     all(predicate, values(::PDVec); kwargs...)
@@ -754,16 +786,16 @@ end
 ### Operator linear algebra operations
 ###
 """
-    mul!(y::PDVec, A::AbstractOperator, x::PDVec[, w::PDWorkingMemory])
+    mul!(y::PDVec, A::AbstractObservable, x::PDVec[, w::PDWorkingMemory])
 
 Perform `y = A * x` in-place. The working memory `w` is required to facilitate
 threaded/distributed operations. If not passed a new instance will be allocated. `y` and `x`
 may be the same vector.
 
-See [`PDVec`](@ref), [`PDWorkingMemory`](@ref), [`AbstractOperator`](@ref).
+See [`PDVec`](@ref), [`PDWorkingMemory`](@ref), [`AbstractObservable`](@ref).
 """
 function LinearAlgebra.mul!(
-    y::PDVec, op::AbstractOperator, x::PDVec,
+    y::PDVec, op::AbstractObservable, x::PDVec,
     w=PDWorkingMemory(y; style=StochasticStyle(y)),
 )
     if !(w.style isa IsDeterministic)
@@ -777,30 +809,30 @@ function LinearAlgebra.mul!(
 end
 
 """
-    dot(y::PDVec, A::AbstractOperator, x::PDVec[, w::PDWorkingMemory])
+    dot(y::PDVec, A::AbstractObservable, x::PDVec[, w::PDWorkingMemory])
 
 Perform `y ⋅ A ⋅ x`. The working memory `w` is required to facilitate threaded/distributed
 operations with non-diagonal `A`. If needed and not passed a new instance will be
 allocated. `A` can be replaced with a tuple of operators.
 
 See [`PDVec`](@ref), [`PDWorkingMemory`](@ref).
 """
-function LinearAlgebra.dot(t::PDVec, op::AbstractOperator, u::PDVec, w)
+function LinearAlgebra.dot(t::PDVec, op::AbstractObservable, u::PDVec, w)
     return dot(LOStructure(op), t, op, u, w)
 end
-function LinearAlgebra.dot(t::PDVec, op::AbstractOperator, u::PDVec)
+function LinearAlgebra.dot(t::PDVec, op::AbstractObservable, u::PDVec)
     return dot(LOStructure(op), t, op, u)
 end
 function LinearAlgebra.dot(
-    ::IsDiagonal, t::PDVec, op::AbstractOperator, u::PDVec, w=nothing
+    ::IsDiagonal, t::PDVec, op::AbstractObservable, u::PDVec, w=nothing
 )
     T = typeof(zero(valtype(t)) * zero(valtype(u)) * zero(eltype(op)))
     return sum(pairs(u); init=zero(T)) do (k, v)
         T(conj(t[k]) * diagonal_element(op, k) * v)
     end
 end
 function LinearAlgebra.dot(
-    ::LOStructure, left::PDVec, op::AbstractOperator, right::PDVec, w=nothing
+    ::LOStructure, left::PDVec, op::AbstractObservable, right::PDVec, w=nothing
 )
     # First two cases: only one vector is distrubuted. Avoid shuffling things around
     # by placing that one on the left to reduce the need for communication.
@@ -819,7 +851,7 @@ function LinearAlgebra.dot(
 end
 # Default variant: also called from other LOStructures.
 function LinearAlgebra.dot(
-    ::AdjointUnknown, t::PDVec, op::AbstractOperator, source::PDVec, w=nothing
+    ::AdjointUnknown, t::PDVec, op::AbstractObservable, source::PDVec, w=nothing
 )
     if is_distributed(t)
         if isnothing(w)
@@ -833,19 +865,6 @@ function LinearAlgebra.dot(
     return dot_from_right(target, op, source)
 end
 
-function Interfaces.dot_from_right(target, op, source::PDVec)
-    T = typeof(zero(valtype(target)) * zero(valtype(source)) * zero(eltype(op)))
-
-    result = sum(pairs(source); init=zero(T)) do (k, v)
-        res = conj(target[k]) * diagonal_element(op, k) * v
-        for (k_off, v_off) in offdiagonals(op, k)
-            res += conj(target[k_off]) * v_off * v
-        end
-        res
-    end
-    return result::T
-end
-
 function LinearAlgebra.dot(t::PDVec, ops::Tuple, source::PDVec, w=nothing)
     if is_distributed(t) && any(LOStructure(op) ≢ IsDiagonal() for op in ops)
         if isnothing(w)

src/Hamiltonians/Hamiltonians.jl

@@ -65,6 +65,7 @@ using TupleTools: TupleTools
 
 using ..BitStringAddresses
 using ..Interfaces
+using ..Interfaces: sum_mutating!
 import ..Interfaces: diagonal_element, num_offdiagonals, get_offdiagonal, starting_address,
     offdiagonals, random_offdiagonal, LOStructure, allows_address_type
 
@@ -85,7 +86,7 @@ export FroehlichPolaron
 export ParticleNumberOperator
 
 export G2MomCorrelator, G2RealCorrelator, G2RealSpace, SuperfluidCorrelator, DensityMatrixDiagonal, Momentum
-export SingleParticleExcitation, TwoParticleExcitation
+export SingleParticleExcitation, TwoParticleExcitation, ReducedDensityMatrix
 export StringCorrelator
 
 export CubicGrid, PeriodicBoundaries, HardwallBoundaries, LadderBoundaries
@@ -94,6 +95,14 @@ export HOCartesianContactInteractions, HOCartesianEnergyConservedPerDim, HOCarte
 export AxialAngularMomentumHO
 export get_all_blocks, fock_to_cart
 
+if VERSION < v"1.10"
+    # used for ReducedDensityMatrix
+    function hermitianpart!(A)
+        A .= (A + A') / 2
+        return Hermitian(A)
+    end
+end
+
 include("abstract.jl")
 include("offdiagonals.jl")
 include("geometry.jl")

src/Hamiltonians/abstract.jl

@@ -1,5 +1,5 @@
 """
-    check_address_type(h::AbstractOperator, addr_or_type)
+    check_address_type(h::AbstractObservable, addr_or_type)
 Throw an `ArgumentError` if `addr_or_type` is not compatible with `h`, otherwise return
 `true`. Acceptable arguments are either an address or an address type, or a tuple or
 array thereof.
@@ -13,18 +13,18 @@ See also [`allows_address_type`](@ref).
     return true
 end
 @inline check_address_type(h, addr) = check_address_type(h, typeof(addr))
-@inline function check_address_type(h::AbstractOperator, v::Union{AbstractArray,Tuple})
+@inline function check_address_type(h::AbstractObservable, v::Union{AbstractArray,Tuple})
     all(check_address_type(h, a) for a in v)
 end
 
-(h::AbstractOperator)(v) = h * v
-(h::AbstractOperator)(w, v) = mul!(w, h, v)
+(h::AbstractObservable)(v) = h * v
+(h::AbstractObservable)(w, v) = mul!(w, h, v)
 
 BitStringAddresses.num_modes(h::AbstractHamiltonian) = num_modes(starting_address(h))
 
 """
     dimension(h::AbstractHamiltonian, addr=starting_address(h))
-    dimension(h::AbstractOperator, addr)
+    dimension(h::AbstractObservable, addr)
     dimension(addr::AbstractFockAddress)
     dimension(::Type{<:AbstractFockAddress})
 
@@ -66,7 +66,7 @@ When extending [`AbstractFockAddress`](@ref), define a method for
 `dimension(::Type{MyNewFockAddress})`.
 """
 dimension(h::AbstractHamiltonian) = dimension(h, starting_address(h))
-dimension(::AbstractOperator, addr) = dimension(addr)
+dimension(::AbstractObservable, addr) = dimension(addr)
 dimension(addr::AbstractFockAddress) = dimension(typeof(addr))
 dimension(::T) where {T<:Number} = typemax(T) # e.g. integer addresses
 
@@ -95,10 +95,10 @@ function dimension(::Type{T}, h, addr=starting_address(h)) where {T}
 end
 dimension(::Type{T}, addr::AbstractFockAddress) where {T} = T(dimension(addr))
 
-Base.isreal(h::AbstractOperator) = eltype(h) <: Real
-LinearAlgebra.isdiag(h::AbstractOperator) = LOStructure(h) ≡ IsDiagonal()
-LinearAlgebra.ishermitian(h::AbstractOperator) = LOStructure(h) ≡ IsHermitian()
-LinearAlgebra.issymmetric(h::AbstractOperator) = ishermitian(h) && isreal(h)
+Base.isreal(h::AbstractObservable) = eltype(h) <: Real
+LinearAlgebra.isdiag(h::AbstractObservable) = LOStructure(h) ≡ IsDiagonal()
+LinearAlgebra.ishermitian(h::AbstractObservable) = LOStructure(h) ≡ IsHermitian()
+LinearAlgebra.issymmetric(h::AbstractObservable) = ishermitian(h) && isreal(h)
 
 BitStringAddresses.near_uniform(h::AbstractHamiltonian) = near_uniform(typeof(starting_address(h)))
 
@@ -211,7 +211,7 @@ Part of the [`AbstractHamiltonian`](@ref) interface.
 """
 momentum
 
-function Base.getindex(ham::AbstractOperator{T}, address1, address2) where {T}
+function Base.getindex(ham::AbstractObservable{T}, address1, address2) where {T}
     # calculate the matrix element when only two bitstring addresses are given
     # this is NOT used for the QMC algorithm and is currently not used either
     # for building the matrix for conventional diagonalisation.
@@ -227,17 +227,17 @@ function Base.getindex(ham::AbstractOperator{T}, address1, address2) where {T}
     return res
 end
 
-LinearAlgebra.adjoint(op::AbstractOperator) = adjoint(LOStructure(op), op)
+LinearAlgebra.adjoint(op::AbstractObservable) = adjoint(LOStructure(op), op)
 
 """
-    adjoint(::LOStructure, op::AbstractOperator)
+    adjoint(::LOStructure, op::AbstractObservable)
 
-Represent the adjoint of an [`AbstractOperator`](@ref). Extend this method to define
+Represent the adjoint of an [`AbstractObservable`](@ref). Extend this method to define
 custom adjoints.
 """
 function LinearAlgebra.adjoint(::S, op) where {S<:LOStructure}
     throw(ArgumentError(
-        "`adjoint()` is not defined for `AbstractOperator`s with `LOStructure` `$(S)`. "*
+        "`adjoint()` is not defined for `AbstractObservable`s with `LOStructure` `$(S)`. "*
         " Is your Hamiltonian hermitian?"
     ))
 end
@@ -252,7 +252,7 @@ function LinearAlgebra.adjoint(::IsDiagonal, op)
 end
 
 """
-    TransformUndoer(transform::AbstractHamiltonian, op::AbstractOperator) <: AbstractHamiltonian
+    TransformUndoer(transform::AbstractHamiltonian, op::AbstractObservable) <: AbstractHamiltonian
 
 Create a new operator for the purpose of calculating overlaps of transformed
 vectors, which are defined by some transformation `transform`. The new operator should
@@ -280,7 +280,7 @@ to represent ``f^{-1} A f^{-1}``.
 * [`GuidingVectorSampling`](@ref)
 """
 struct TransformUndoer{
-    T,K<:AbstractHamiltonian,O<:Union{AbstractOperator,Nothing}
+    T,K<:AbstractHamiltonian,O<:Union{AbstractObservable,Nothing}
 } <: AbstractHamiltonian{T}
     transform::K
     op::O