Excited states

src/ExactDiagonalization/basis_set_representation.jl

@@ -95,7 +95,7 @@ end
 """
     build_basis(
         ham, address=starting_address(ham);
-        cutoff, filter, sizelim, sort=false, kwargs...
+        cutoff, filter, sizelim, sort=false, minimum_size, kwargs...
     ) -> basis
     build_basis(ham, addresses::AbstractVector; kwargs...)
 
@@ -109,7 +109,10 @@ than `cutoff`. Alternatively, an arbitrary `filter` function can be used instead
 Addresses passed as arguments are not filtered.
 A maximum basis size `sizelim` can be set which will throw an error if the expected
 dimension of `ham` is larger than `sizelim`. This may be useful when memory may be a
-concern. These options are disabled by default.
+concern.
+Setting a `minimum_size` will stop generating addresses once at least `minimum_size`
+addresses have been generated, rather than returning the full basis.
+These options are disabled by default.
 
 Setting `sort` to `true` will sort the basis. Any additional keyword arguments
 are passed on to `Base.sort!`.
@@ -121,19 +124,20 @@ function build_basis(
     sort=false,
     max_size=Inf, # retained for backwards compatibility; use sizelim instead
     sizelim=max_size,
+    minimum_size=Inf,
     kwargs...
 )
     check_address_type(ham, addr_or_vec)
     single_addr = addr_or_vec isa Union{AbstractArray,Tuple} ? addr_or_vec[1] : addr_or_vec
-    if dimension(ham, single_addr) > sizelim
+    if minimum([dimension(ham, single_addr), minimum_size]) > sizelim
         throw(ArgumentError("dimension larger than sizelim"))
     end
     # Set up `adds` as queue of addresses. Also returned as the basis.
     adds = addr_or_vec isa Union{AbstractArray,Tuple} ? [addr_or_vec...] : [addr_or_vec]
     known_basis = Set(adds)     # known addresses
 
     i = 0
-    while i < length(adds)
+    while i < length(adds) < minimum_size
         i += 1
         add = adds[i]
 

src/Rimu.jl

@@ -3,7 +3,7 @@ module Rimu
 using Arrow: Arrow
 using DataFrames: DataFrames, DataFrame, metadata
 using DataStructures: DataStructures
-using LinearAlgebra: LinearAlgebra, dot, isdiag
+using LinearAlgebra: LinearAlgebra, dot, isdiag, eigvecs, norm
 using OrderedCollections: OrderedCollections, LittleDict, freeze
 using Parameters: Parameters, @pack!, @unpack, @with_kw
 using ProgressLogging: ProgressLogging, @logprogress, @withprogress

src/fciqmc.jl

@@ -28,8 +28,8 @@ end
 Set up the initial shift parameters for the [`FCIQMC`](@ref) algorithm.
 """
 function set_up_initial_shift_parameters(algorithm::FCIQMC, hamiltonian,
-    starting_vectors::SMatrix{S,R}, shift, time_step
-) where {S,R}
+    starting_vectors::SMatrix{R,S}, shift, time_step
+) where {R,S}
     shift_strategy = algorithm.shift_strategy
 
     if shift === nothing
@@ -45,7 +45,7 @@ function set_up_initial_shift_parameters(algorithm::FCIQMC, hamiltonian,
         initialise_shift_parameters(shift_strategy, s, walkernumber(sv), time_step)
     end)
     @assert length(initial_shift_parameters) == S * R
-    return SMatrix{S,R}(initial_shift_parameters)
+    return SMatrix{R,S}(initial_shift_parameters)
 end
 
 function _determine_initial_shift(hamiltonian, starting_vectors)
@@ -192,4 +192,20 @@ end
 function advance!(algorithm, report, state::ReplicaState, replica::SpectralState{1})
     return advance!(algorithm, report, state, only(replica.single_states))
 end
-# TODO: add advance! for SpectralState{N} where N > 1
+
+function advance!(algorithm, report, state::ReplicaState, replica::SpectralState{N, <:Any, GramSchmidt{N}}) where {N}
+    proceed = true
+    if state.step[] % replica.spectral_strategy.orthogonalization_interval == 0
+        for i in 1:N
+            for j in 1:i-1
+                u = replica[i].v
+                v = replica[j].v
+                add!(u, v, -dot(u, v) / norm(v)^2)
+            end
+        end
+    end
+    for i in 1:N
+        proceed &= advance!(algorithm, report, state, replica[i])
+    end
+    return proceed
+end

src/pmc_simulation.jl

@@ -21,42 +21,57 @@ mutable struct PMCSimulation
 end
 
 function _set_up_starting_vectors(
-    start_at, n_replicas, n_spectral, style, initiator, threading, copy_vectors
+    ham, start_at, n_replicas, n_spectral, style, initiator, threading, copy_vectors
 )
     if start_at isa AbstractMatrix && size(start_at) == (n_replicas, n_spectral)
         if eltype(start_at) <: Union{AbstractDVec, RMPI.MPIData} # already dvecs
             if copy_vectors
-                return SMatrix{n_spectral, n_replicas}(deepcopy(v) for v in start_at)
+                return SMatrix{n_replicas, n_spectral}(deepcopy(v) for v in start_at)
             else
-                return SMatrix{n_spectral, n_replicas}(v for v in start_at)
+                return SMatrix{n_replicas, n_spectral}(v for v in start_at)
             end
         else
-            return SMatrix{n_spectral, n_replicas}(
+            return SMatrix{n_replicas, n_spectral}(
                 default_starting_vector(sa; style, initiator, threading) for sa in start_at
             )
         end
+    elseif n_spectral > 1
+        basis = build_basis(ham; minimum_size=n_spectral*5)
+        basis = sort!(basis; by=x -> diagonal_element(ham, x))
+        trunc = BasisSetRepresentation(ham, basis; filter=Returns(false), sizelim=Inf)
+        vecs = eigvecs(Matrix(trunc))
+        if threading
+            eigs = [PDVec(zip(basis, 10*vecs[:,i]); style, initiator) for i in 1:n_spectral]
+        else
+            if initiator == NonInitiator()
+                eigs = [DVec(zip(basis, 10*vecs[:,i]); style) for i in 1:n_spectral]
+            else
+                eigs = [InitiatorDVec(zip(basis,10*vecs[:,i]);style,initiator) for i in 1:n_spectral]
+            end
+        end
+        mat = stack([eigs for _ in 1:n_replicas], dims=1)
+        return SMatrix{n_replicas, n_spectral}(mat)
     end
-    n_spectral == 1 || throw(ArgumentError("The number of spectral states must match the number of starting vectors."))
     if start_at isa AbstractVector && length(start_at) == n_replicas
         if eltype(start_at) <: Union{AbstractDVec, RMPI.MPIData} # already dvecs
             if copy_vectors
-                return SMatrix{n_spectral, n_replicas}(deepcopy(v) for v in start_at)
+                return SMatrix{n_replicas, n_spectral}(deepcopy(v) for v in start_at)
             else
-                return SMatrix{n_spectral, n_replicas}(v for v in start_at)
+                return SMatrix{n_replicas, n_spectral}(v for v in start_at)
             end
         else
-            return SMatrix{n_spectral, n_replicas}(
+            return SMatrix{n_replicas, n_spectral}(
                 default_starting_vector(sa; style, initiator, threading) for sa in start_at
             )
         end
     elseif start_at isa Union{AbstractDVec, RMPI.MPIData}
         if n_replicas == 1 && !copy_vectors
-            return SMatrix{n_spectral, n_replicas}(start_at)
+            return SMatrix{n_replicas, n_spectral}(start_at)
         else
-            return SMatrix{n_spectral, n_replicas}(deepcopy(start_at) for _ in 1:n_replicas)
+            return SMatrix{n_replicas, n_spectral}(deepcopy(start_at) for _ in 1:n_replicas)
         end
     end
-    return SMatrix{n_spectral,n_replicas}(
+    return SMatrix{n_replicas, n_spectral}(
         default_starting_vector(start_at; style, initiator, threading) for _ in 1:n_replicas
     )
 end
@@ -80,10 +95,10 @@ function PMCSimulation(problem::ProjectorMonteCarloProblem; copy_vectors=true)
 
     start_at = isnothing(start_at) ? starting_address(hamiltonian) : start_at
     vectors = _set_up_starting_vectors(
-        start_at, n_replicas, n_spectral, style, initiator,
+        hamiltonian, start_at, n_replicas, n_spectral, style, initiator,
         threading, copy_vectors
     )
-    @assert vectors isa SMatrix{n_spectral,n_replicas}
+    @assert vectors isa SMatrix{n_replicas, n_spectral}
 
     # set up initial_shift_parameters
     shift, time_step = nothing, nothing
@@ -93,17 +108,17 @@ function PMCSimulation(problem::ProjectorMonteCarloProblem; copy_vectors=true)
         @unpack shift, time_step = initial_shift_parameters
         set_up_initial_shift_parameters(algorithm, hamiltonian, vectors, shift, time_step)
     else
-        SMatrix{n_spectral,n_replicas}(initial_shift_parameters...)
+        SMatrix{n_replicas, n_spectral}(initial_shift_parameters...)
     end
-    @assert shift_parameters isa SMatrix{n_spectral,n_replicas}
+    @assert shift_parameters isa SMatrix{n_replicas, n_spectral}
 
     # set up the spectral_states
     wm = working_memory(vectors[1, 1])
     spectral_states = ntuple(n_replicas) do i
         SpectralState(
             ntuple(n_spectral) do j
-                v = vectors[j, i]
-                sp = shift_parameters[j, i]
+                v = vectors[i, j]
+                sp = shift_parameters[i, j]
                 id = if n_replicas * n_spectral == 1
                     ""
                 elseif n_spectral == 1

src/projector_monte_carlo_problem.jl

@@ -52,6 +52,8 @@ the [`FCIQMC`](@ref) algorithm.
 - `n_replicas = 1`: Number of synchronised independent simulations.
 - `replica_strategy = NoStats(n_replicas)`: Which results to report from replica
   simulations, see [`ReplicaStrategy`](@ref).
+- `spectral_strategy = GramSchmidt()`: The [`SpectralStrategy`](@ref) used for simulations
+  of spectral states.
 
 # Example
 

src/strategies_and_params/spectralstrategy.jl

@@ -4,6 +4,10 @@
 
 Abstract type for spectral strategies. The spectral strategy is used to control the number
 of spectral states used in the simulation.
+
+## Implemented Strategies
+
+* [`GramSchmidt`](@ref): Orthogonalize the spectral states using the Gram-Schmidt procedure.
 """
 abstract type SpectralStrategy{S} end
 
@@ -15,14 +19,22 @@ Return the number of spectral states used in the simulation.
 num_spectral_states(::SpectralStrategy{S}) where {S} = S
 
 """
-    GramSchmidt{S} <: SpectralStrategy{S}
-
+    GramSchmidt(S; orthogonalization_interval = 1) <: SpectralStrategy{S} 
 Use the Gram-Schmidt procedure to orthogonalize the excited states. A total of `S` spectral
-states are used in the simulation.
+states are used in the simulation, and they are orthogonalized every 
+`orthogonalization_interval` steps.
+
+Use with the keyword argument `spectral_strategy` in [`ProjectorMonteCarloProblem`](@ref).
 """
-struct GramSchmidt{S} <: SpectralStrategy{S} end
+struct GramSchmidt{S} <: SpectralStrategy{S}
+    orthogonalization_interval::Int
+end
+
+function GramSchmidt(S = 1; orthogonalization_interval = 1)
+    return GramSchmidt{S}(orthogonalization_interval)
+end
 
-function GramSchmidt(num_spectral_states = 1)
-    num_spectral_states > 1 && throw(ArgumentError("num_spectral_states > 1 is currently not supported."))
-    GramSchmidt{num_spectral_states}()
+function Base.show(io::IO, gs::GramSchmidt{S}) where {S}
+    print(io, "GramSchmidt(", S, "; orthogonalization_interval=")
+    print(io, gs.orthogonalization_interval, ")")
 end

test/excited_states_tests.jl

@@ -0,0 +1,41 @@
+using Rimu
+using Test
+
+
+@testset "excited state energies" begin
+    ham = HubbardReal1D(BoseFS(1,1,1,1,1))
+    pr = ExactDiagonalizationProblem(ham)
+    vals = solve(pr).values
+
+    spectral_strategy = GramSchmidt(3)
+    last_step=2000
+    style = IsDeterministic()
+    p = ProjectorMonteCarloProblem(ham; spectral_strategy, last_step, style)
+    s = solve(p)
+    df = DataFrame(s)
+    energy1 = shift_estimator(df, shift="shift_s1_1", skip=1000)
+    energy2 = shift_estimator(df, shift="shift_s2_1", skip=1000)
+    energy3 = shift_estimator(df, shift="shift_s3_1", skip=1000)
+
+    @test energy1.mean ≈ vals[1]
+    @test energy2.mean ≈ vals[2]
+    @test energy3.mean ≈ vals[3]
+
+    n_replicas = 2
+    p = ProjectorMonteCarloProblem(ham; spectral_strategy, last_step, style, n_replicas)
+    s = solve(p)
+    df = DataFrame(s)
+    energy1 = shift_estimator(df, shift="shift_s1_1", skip=1000)
+    energy2 = shift_estimator(df, shift="shift_s2_1", skip=1000)
+    energy3 = shift_estimator(df, shift="shift_s3_1", skip=1000)
+    energy4 = shift_estimator(df, shift="shift_s1_2", skip=1000)
+    energy5 = shift_estimator(df, shift="shift_s2_2", skip=1000)
+    energy6 = shift_estimator(df, shift="shift_s3_2", skip=1000)
+
+    @test energy1.mean ≈ vals[1]
+    @test energy2.mean ≈ vals[2]
+    @test energy3.mean ≈ vals[3]
+    @test energy4.mean ≈ vals[1]
+    @test energy5.mean ≈ vals[2]
+    @test energy6.mean ≈ vals[3]
+end

test/runtests.jl

@@ -183,4 +183,8 @@ end
     # For example, replace `using Foo` with `using Foo: bar, baz`.
 end
 
+@safetestset "excited states" begin
+    include("excited_states_tests.jl")
+end
+
 # Note: Running Rimu with several MPI ranks is tested seperately on GitHub CI and not here.