export num_replicas and num_spectral_states; new way of setting up st…

…arting vectors

src/PMCSimulation.jl

@@ -21,76 +21,100 @@ mutable struct PMCSimulation
     elapsed_time::Float64
 end
 
-function _set_up_v(dv::AbstractDVec, copy_vectors, _...)
-    # we are not changing type, style or initiator because an AbstractDVec is passed
-    if copy_vectors
-        return deepcopy(dv)
-    else
-        return dv
+function _set_up_starting_vectors(
+    start_at, n_replicas, n_spectral, style, initiator, threading, copy_vectors
+)
+    if start_at isa AbstractMatrix && size(start_at) == (n_spectral, 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)
+            else
+                return SMatrix{n_spectral, n_replicas}(v for v in start_at)
+            end
+        else
+            return SMatrix{n_spectral, n_replicas}(
+                default_starting_vector(sa; style, initiator, threading) for sa in start_at
+            )
+        end
     end
-end
-
-function _set_up_v(fdv::FrozenDVec, _, style, initiator, threading)
-    # we are allocating new memory
-    if threading
-        v = PDVec(fdv; style, initiator)
-    elseif initiator isa NonInitiator
-        v = DVec(fdv; style)
-    else
-        v = InitiatorDVec(fdv; style, initiator)
+    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)
+            else
+                return SMatrix{n_spectral, n_replicas}(v for v in start_at)
+            end
+        else
+            return SMatrix{n_spectral, n_replicas}(
+                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)
+        else
+            return SMatrix{n_spectral, n_replicas}(deepcopy(start_at) for _ in 1:n_replicas)
+        end
     end
-    return v
+    return SMatrix{n_spectral,n_replicas}(
+        default_starting_vector(start_at; style, initiator, threading) for _ in 1:n_replicas
+    )
 end
 
 function PMCSimulation(problem::ProjectorMonteCarloProblem; copy_vectors=true)
-    @unpack algorithm, hamiltonian, starting_vectors, style, threading, simulation_plan,
+    # @unpack algorithm, hamiltonian, starting_vectors, style, threading, simulation_plan,
+    @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
 
     reporting_strategy = refine_reporting_strategy(reporting_strategy)
 
     n_replicas = num_replicas(replica_strategy)
+    n_spectral = num_spectral_states(spectral_strategy)
 
     # seed the random number generator
     if !isnothing(random_seed)
         Random.seed!(random_seed + hash(RMPI.mpi_rank()))
     end
 
-    # set up the starting vectors and shift parameters
-    if length(starting_vectors) == 1 && n_replicas > 1
-        # in this case we always make new copies of the starting vector
-        vectors = ntuple(n_replicas) do i
-            v = _set_up_v(only(starting_vectors), true, style, initiator, threading)
-            return v
-        end
-        shift_parameters = ntuple(n_replicas) do i
-            deepcopy(only(initial_shift_parameters))
-        end
+    start_at = isnothing(start_at) ? starting_address(hamiltonian) : start_at
+    vectors = _set_up_starting_vectors(
+        start_at, n_replicas, n_spectral, style, initiator,
+        threading, copy_vectors
+    )
+    @assert vectors isa SMatrix{n_spectral,n_replicas}
+
+    # set up initial_shift_parameters
+    shift, time_step = nothing, nothing
+
+    shift_parameters = if initial_shift_parameters isa Union{NamedTuple, DefaultShiftParameters}
+        # we have previously stored shift and time_step
+        @unpack shift, time_step = initial_shift_parameters
+        set_up_initial_shift_parameters(algorithm, hamiltonian, vectors, shift, time_step)
     else
-        @assert length(starting_vectors) == n_replicas == length(initial_shift_parameters)
-        vectors = map(starting_vectors) do dv
-            v = _set_up_v(dv, copy_vectors, style, initiator, threading)
-            return v
-        end
-        shift_parameters = deepcopy(initial_shift_parameters)
+        initial_shift_parameters
     end
+    @assert shift_parameters isa SMatrix{n_spectral,n_replicas}
 
     # set up the spectral_states
     spectral_states = ntuple(n_replicas) do i
-        v, sp = vectors[i], shift_parameters[i]
         SpectralState(
-            (SingleState(
-                hamiltonian,
-                algorithm,
-                v,
-                zerovector(v),
-                working_memory(v),
-                sp,
-                n_replicas == 1 ? "" : "_$i"
-            ),),
-            spectral_strategy
-        )
+            ntuple(n_spectral) do j
+                v = vectors[j, i]
+                sp = shift_parameters[j, i]
+                id = if n_replicas * n_spectral == 1
+                    ""
+                elseif n_spectral == 1
+                    "_$(i)"
+                else
+                    "_s$(j)_$(i)" # j is the spectral state index, i is the replica index
+                end # we have to think about how to label the spectral states
+                SingleState(
+                    hamiltonian, algorithm, v, zerovector(v), working_memory(v), sp, id
+                )
+            end, spectral_strategy)
     end
     @assert spectral_states isa NTuple{n_replicas, <:SpectralState}
 

src/ProjectorMonteCarloProblem.jl

@@ -28,11 +28,9 @@ the [`FCIQMC`](@ref) algorithm.
 - `time_step = 0.01`: Initial time step size.
 - `last_step = 100`: Controls the number of steps.
 - `targetwalkers = 1_000`: Target for the 1-norm of the coefficient vector.
-- `start_at = starting_address(hamiltonian)`: Define the initial state vector. This can be a
-    single address, a collection of addresses, a single starting vector, or a collection of
-    starting vectors. If multiple starting vectors are passed, their number must match the
-    number of replicas. If (a collection of) [`AbstractDVec`](@ref)(s) is passed, the
-    keyword arguments `style`, `initiator`, and `threading` are ignored.
+- `start_at = starting_address(hamiltonian)`: Define the initial state vector(s).
+    If multiple starting vectors are passed, their number must match the product of the
+    number of replicas and the number of spectral states.
 - `style = IsDynamicSemistochastic()`: The [`StochasticStyle`](@ref) of the simulation.
 - `initiator = NonInitiator()`: Whether to use initiators. See [`InitiatorRule`](@ref).
 - `threading`: Default is to use multithreading and
@@ -93,13 +91,13 @@ struct ProjectorMonteCarloProblem{N,S} # is not type stable but does not matter
     # N is the number of replicas, S is the number of spectral states
     algorithm::PMCAlgorithm
     hamiltonian::AbstractHamiltonian
-    starting_vectors
+    start_at  # starting_vectors
     style::StochasticStyle
     initiator::InitiatorRule
     threading::Bool
     simulation_plan::SimulationPlan
     replica_strategy::ReplicaStrategy{N}
-    initial_shift_parameters::Tuple
+    initial_shift_parameters
     reporting_strategy::ReportingStrategy
     post_step_strategy::Tuple
     spectral_strategy::SpectralStrategy{S}
@@ -167,31 +165,12 @@ function ProjectorMonteCarloProblem(
         random_seed = UInt64(random_seed)
     end
 
-    # set up starting_vectors
-    if start_at isa AbstractFockAddress # single address
-        starting_vectors = (freeze(DVec(start_at => 1,)),) # tuple of length 1
-    elseif eltype(start_at) <: AbstractFockAddress # multiple addresses
-        starting_vectors = (freeze(DVec(address => 1 for address in start_at)),)
-    elseif start_at isa Union{AbstractDVec, RMPI.MPIData} # single starting vector
-        starting_vectors = (start_at,) # tuple of length 1
-    elseif eltype(start_at) <: Pair{<:AbstractFockAddress} # single starting vector
-        starting_vectors = (freeze(DVec(start_at)),) # tuple of length 1
-    elseif eltype(start_at) <: AbstractDVec # multiple starting vectors
-        starting_vectors = Tuple(sv for sv in start_at)
-    else
-        throw(ArgumentError("`start_at` has invalid format."))
+    # a proper setup of initial_shift_parameters is done in PMCSimulation
+    # here we just store the initial shift and time_step if initial_shift_parameters is not
+    # provided
+    if isnothing(initial_shift_parameters)
+        initial_shift_parameters = (; shift, time_step)
     end
-    @assert starting_vectors isa NTuple{<:Any,<:Union{AbstractDVec,FrozenDVec,RMPI.MPIData}}
-    length(starting_vectors) == 1 || length(starting_vectors) == n_replicas ||
-        throw(ArgumentError("The number of starting vectors must match the number of replicas."))
-    all(v -> keytype(v) <: allowed_address_type(hamiltonian), starting_vectors) ||
-        throw(ArgumentError("The address type is not allowed for the Hamiltonian."))
-
-    # set up initial_shift_parameters
-    initial_shift_parameters = set_up_initial_shift_parameters(algorithm, hamiltonian,
-        starting_vectors, shift, time_step, initial_shift_parameters)
-
-    @assert length(initial_shift_parameters) == length(starting_vectors)
 
     shift_strategy = algorithm.shift_strategy
     if isnothing(maxlength)
@@ -213,7 +192,7 @@ function ProjectorMonteCarloProblem(
     return ProjectorMonteCarloProblem{n_replicas,num_spectral_states(spectral_strategy)}(
         algorithm,
         hamiltonian,
-        starting_vectors,
+        start_at, # starting_vectors,
         style,
         initiator,
         threading,

src/RMPI/mpidata.jl

@@ -66,15 +66,6 @@ function Base.similar(md::MPIData)
     return MPIData(similar(md.data), md.comm, md.root, md.s)
 end
 
-function Rimu._set_up_v(dv::MPIData, copy_vectors, _...)
-    # we are not changing type, style or initiator because a wrapped AbstractDVec is passed
-    if copy_vectors
-        return deepcopy(dv)
-    else
-        return dv
-    end
-end
-
 ###
 ### Iterators
 ###

src/Rimu.jl

@@ -73,7 +73,7 @@ export TimeStepStrategy, ConstantTimeStep, OvershootControl
 export localpart, walkernumber
 export smart_logger, default_logger
 export ProjectorMonteCarloProblem, SimulationPlan, state_vectors, single_states
-export FCIQMC
+export FCIQMC, num_replicas, num_spectral_states
 
 function __init__()
     # Turn on smart logging once at runtime. Turn off with `default_logger()`.

src/fciqmc.jl

@@ -25,29 +25,27 @@ end
     starting_vectors, shift, time_step, initial_shift_parameters
 )
 
-Set up the initial shift parameters for the FCIQMC algorithm.
+Set up the initial shift parameters for the [`FCIQMC`](@ref) algorithm.
 """
 function set_up_initial_shift_parameters(algorithm::FCIQMC, hamiltonian,
-    starting_vectors, shift, time_step, initial_shift_parameters
-)
+    starting_vectors::SMatrix{S,R}, shift, time_step
+) where {S,R}
     shift_strategy = algorithm.shift_strategy
-    if isnothing(initial_shift_parameters)
-        if shift === nothing
-            initial_shifts = _determine_initial_shift(hamiltonian, starting_vectors)
-        elseif shift isa Number
-            initial_shifts = [float(shift) for _ in 1:length(starting_vectors)]
-        elseif length(shift) == length(starting_vectors)
-            initial_shifts = float.(shift)
-        else
-            throw(ArgumentError("The number of shifts must match the number of starting vectors."))
-        end
-        initial_shift_parameters = Tuple(map(zip(starting_vectors, initial_shifts)) do (sv, s)
-            initialise_shift_parameters(shift_strategy, s, walkernumber(sv), time_step)
-        end)
-    elseif !(initial_shift_parameters isa Tuple)
-        initial_shift_parameters = Tuple(initial_shift_parameters for _ in 1:length(starting_vectors))
+
+    if shift === nothing
+        initial_shifts = _determine_initial_shift(hamiltonian, starting_vectors)
+    elseif shift isa Number
+        initial_shifts = [float(shift) for _ in 1 : S*R] # not great for excited states
+    elseif length(shift) == n_spectral
+        initial_shifts = float.(shift)
+    else
+        throw(ArgumentError("The number of shifts must match the number of starting vectors."))
     end
-    return initial_shift_parameters
+    initial_shift_parameters = Tuple(map(zip(starting_vectors, initial_shifts)) do (sv, s)
+        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)
 end
 
 function _determine_initial_shift(hamiltonian, starting_vectors)

src/lomc.jl

@@ -178,23 +178,6 @@ function lomc!(
     end
 end
 
-
-
-#     state = ReplicaState(
-#         ham, v;
-#         reporting_strategy = r_strat, # deprecations
-#         shift_strategy = s_strat,
-#         time_step_strategy = τ_strat,
-#         replica_strategy = replica,
-#         post_step_strategy = post_step,
-#         starting_step,
-#         last_step,
-#         time_step,
-#         shift,
-#         kwargs...
-#     )
-#     return lomc!(state, df; name, metadata)
-# end
 function lomc!(
     ham;
     style=IsStochasticInteger(),
@@ -211,24 +194,6 @@ function lomc!(::AbstractMatrix, v=nothing; kwargs...)
     throw(ArgumentError("Using lomc! with a matrix is no longer supported. Use `MatrixHamiltonian` instead."))
 end
 
-# # continuation run
-# function lomc!(sm::PMCSimulation, df=DataFrame(); laststep=0, name="lomc!", kwargs...)
-#     solve!(sm; last_step=laststep, display_name=name, kwargs...)
-
-#     # Put report into DataFrame and merge with `df`. We are assuming the previous `df` is
-#     # compatible, which should be the case if the run is an actual continuation. Maybe the
-#     # DataFrames should be merged in a more permissive manner?
-#     result_df = DataFrame(sm)
-
-#     if !isempty(df)
-#         result_df = vcat(df, result_df) # metadata is not propagated
-#         for (key, val) in get_metadata(report) # add metadata
-#             DataFrames.metadata!(result_df, key, val)
-#         end
-#     end
-#     return (; df=result_df, state=sm)
-# end
-
 # methods for backward compatibility
 function lomc!(state::ReplicaState, df=DataFrame(); laststep=0, name="lomc!", metadata=nothing)
     if !iszero(laststep)

src/qmc_states.jl

@@ -317,20 +317,33 @@ function default_starting_vector(
 )
     return default_starting_vector(address; kwargs...)
 end
-function default_starting_vector(address;
-    style=IsStochasticInteger(),
-    threading=nothing,
-    initiator=NonInitiator(),
+
+function default_starting_vector(address::AbstractFockAddress; population=10, kwargs...)
+    return default_starting_vector(address=>population; kwargs...)
+end
+
+function default_starting_vector(fdv::Union{FrozenDVec,Pair};
+    style = IsStochasticInteger(),
+    threading = nothing,
+    initiator = NonInitiator(),
+    mpi = RMPI.mpi_size() > 1,
 )
-    if isnothing(threading)
-        threading = Threads.nthreads() > 1
-    end
+    threading === nothing && (threading = Threads.nthreads() > 1)
+    return _setup_dvec(fdv, style, initiator, threading, mpi)
+end
+
+function _setup_dvec(fdv::Union{FrozenDVec,Pair}, style, initiator, threading, mpi)
+    # we are allocating new memory
     if threading
-        v = PDVec(address => 10; style, initiator)
-    elseif initiator isa NonInitiator
-        v = DVec(address => 10; style)
+        return PDVec(fdv; style, initiator)
+    end
+    if initiator isa NonInitiator
+        v = DVec(fdv; style)
     else
-        v = InitiatorDVec(address => 10; style, initiator)
+        v = InitiatorDVec(fdv; style, initiator)
+    end
+    if mpi
+        return RMPI.MPIData(v)
     end
     return v
 end