Feature/reduced density matrix

Project.toml

@@ -1,7 +1,7 @@
 name = "Rimu"
 uuid = "c53c40cc-bd84-11e9-2cf4-a9fde2b9386e"
 authors = ["Joachim Brand <j.brand@massey.ac.nz>"]
-version = "0.13.0"
+version = "0.13.2-dev"
 
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"

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

@@ -70,15 +70,18 @@ Stoquastic
 ```
 
 ## Observables
-Observables are [`AbstractOperator`](@ref)s that represent a physical
-observable. Their expectation values can be sampled during a
-[`ProjectorMonteCarloProblem`](@ref) simulation by passing
-them into a suitable [`ReplicaStrategy`](@ref), e.g. 
-[`AllOverlaps`](@ref).  [`AbstractOperator`](@ref) is a supertype of 
-[`AbstractHamiltonian`](@ref) and has less stringent 
-requirements. Some observables are also [`AbstractHamiltonian`](@ref)s.
+`Rimu.jl` offers two other supertypes for operators that are less 
+restrictive than [`AbstractHamiltonian`](@ref). 
+[`AbstractObservable`](@ref) and [`AbstractOperator`](@ref)s both
+can represent a physical observable. Their expectation values can be sampled during a [`ProjectorMonteCarloProblem`](@ref) simulation by 
+passing them into a suitable [`ReplicaStrategy`](@ref), e.g. 
+[`AllOverlaps`](@ref). Some observables are also [`AbstractHamiltonian`](@ref)s. The full type hierarchy is
+```julia
+AbstractHamiltonian{T} <: AbstractOperator{T} <: AbstractObservable{T}
+```
 
 ```@docs
+AbstractObservable
 AbstractOperator
 ParticleNumberOperator
 G2RealCorrelator
@@ -89,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.
@@ -531,6 +531,38 @@ function Base.mapreduce(f::F, op::O, t::PDVecIterator; kwargs...) where {F,O}
     return merge_remote_reductions(t.vector.communicator, op, result)
 end
 
+"""
+    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_non_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_non_mutating(accu, f!, iterator::PDVecIterator; kwargs...)
+    interim_result = Folds.mapreduce(
+        +, 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)
+end
+
 """
     all(predicate, keys(::PDVec); kwargs...)
     all(predicate, values(::PDVec); kwargs...)
@@ -748,16 +780,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)
@@ -771,30 +803,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.
@@ -813,7 +845,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)
@@ -827,19 +859,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")