Excited states

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
 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) / dot(v, v))
+ 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; cutoff=diagonal_element(ham, starting_address(ham)) + 5)
+ basis = sort!(basis; by=x -> diagonal_element(ham, x))[1:2*n_spectral]
+ 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/strategies_and_params/spectralstrategy.jl

@@ -15,14 +15,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.