From c8ffb31c1049b90b68f11f7e7e0aa500670bbd1d Mon Sep 17 00:00:00 2001
From: jamie-tay <jamiet96@gmail.com>
Date: Sun, 24 Nov 2024 14:14:10 +1300
Subject: [PATCH] Add minimum_size keyword argument to
 ProjectorMonteCarloProblem

---
 src/pmc_simulation.jl                | 9 ++++-----
 src/projector_monte_carlo_problem.jl | 8 +++++++-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/pmc_simulation.jl b/src/pmc_simulation.jl
index e8a476964..872d0fb07 100644
--- a/src/pmc_simulation.jl
+++ b/src/pmc_simulation.jl
@@ -21,7 +21,7 @@ mutable struct PMCSimulation
 end
 
 function _set_up_starting_vectors(
-    ham, start_at, n_replicas, n_spectral, style, initiator, threading, copy_vectors
+    ham, start_at, n_replicas, n_spectral, style, initiator, threading, copy_vectors, minimum_size
 )
     if start_at isa AbstractMatrix && size(start_at) == (n_replicas, n_spectral)
         if eltype(start_at) <: Union{AbstractDVec, RMPI.MPIData} # already dvecs
@@ -36,8 +36,7 @@ function _set_up_starting_vectors(
             )
         end
     elseif n_spectral > 1
-        basis = build_basis(ham; minimum_size=n_spectral*5)
-        basis = sort!(basis; by=x -> diagonal_element(ham, x))
+        basis = build_basis(ham; minimum_size)
         trunc = BasisSetRepresentation(ham, basis; filter=Returns(false), sizelim=Inf)
         vecs = eigvecs(Matrix(trunc))
         if threading
@@ -81,7 +80,7 @@ function PMCSimulation(problem::ProjectorMonteCarloProblem; copy_vectors=true)
     @unpack algorithm, hamiltonian, start_at, style, threading, simulation_plan,
         replica_strategy, initial_shift_parameters,
         reporting_strategy, post_step_strategy,
-        maxlength, metadata, initiator, random_seed, spectral_strategy = problem
+        maxlength, metadata, initiator, random_seed, spectral_strategy, minimum_size = problem
 
     reporting_strategy = refine_reporting_strategy(reporting_strategy)
 
@@ -96,7 +95,7 @@ function PMCSimulation(problem::ProjectorMonteCarloProblem; copy_vectors=true)
     start_at = isnothing(start_at) ? starting_address(hamiltonian) : start_at
     vectors = _set_up_starting_vectors(
         hamiltonian, start_at, n_replicas, n_spectral, style, initiator,
-        threading, copy_vectors
+        threading, copy_vectors, minimum_size
     )
     @assert vectors isa SMatrix{n_replicas, n_spectral}
 
diff --git a/src/projector_monte_carlo_problem.jl b/src/projector_monte_carlo_problem.jl
index d1632ca8b..6aa2ee61b 100644
--- a/src/projector_monte_carlo_problem.jl
+++ b/src/projector_monte_carlo_problem.jl
@@ -97,6 +97,9 @@ julia> size(DataFrame(simulation))
     `true`, a random seed is generated. If set to number, this number is used as the seed.
     The seed is used by `solve` such that `solve`ing the problem twice will yield identical
     results. If set to `false`, no seed is used and results are not reproducible.
+- `minimum_size = 2*num_spectral_states(spectral_strategy)`: The minimum size of the basis
+    used to construct starting vectors for simulations of spectral states, if `start_at`
+    is not provided.
 
 See also [`init`](@ref), [`solve`](@ref).
 """
@@ -117,6 +120,7 @@ struct ProjectorMonteCarloProblem{N,S} # is not type stable but does not matter
     maxlength::Int
     metadata::LittleDict{String,String} # user-supplied metadata + display_name
     random_seed::Union{Nothing,UInt64}
+    minimum_size::Int
 end
 
 function Base.show(io::IO, p::ProjectorMonteCarloProblem)
@@ -164,6 +168,7 @@ function ProjectorMonteCarloProblem(
     reporting_strategy = ReportDFAndInfo(),
     post_step_strategy = (),
     spectral_strategy = GramSchmidt(),
+    minimum_size = 2*num_spectral_states(spectral_strategy),
     maxlength = nothing,
     metadata = nothing,
     display_name = "PMCSimulation",
@@ -236,7 +241,8 @@ function ProjectorMonteCarloProblem(
         spectral_strategy,
         maxlength,
         metadata,
-        random_seed
+        random_seed,
+        minimum_size
     )
 end