Restart mechanism with a wall-time budget

Project.toml

@@ -11,6 +11,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
@@ -29,6 +30,7 @@ Distributions = "0.25"
 LinearAlgebra = "1.9"
 Pkg = "1.9"
 Printf = "1.9"
+Serialization = "1.11.0"
 StaticArrays = "1.9"
 Statistics = "1.9"
 TOML = "1"

src/IO/IO.jl

@@ -40,7 +40,7 @@ end
 
 function load_configuration(io, format::Arianna.Format; m=1)
     data = readlines(io)
-    N, box, column_info, metadata = read_header(data, format)
+    t, N, box, column_info, metadata = read_header(data, format) # add t_start
     selrow = get_selrow(format, N, m)
     frame = data[selrow:selrow+N-1]
     bool_molecule = "molecule" in keys(column_info)
@@ -90,7 +90,8 @@ function load_configuration(io, format::Arianna.Format; m=1)
         :box => box,
         :species => species,
         :position => position,
-        :metadata => metadata
+        :metadata => metadata,
+        :t => t
     )
     if bool_molecule
         config_dict[:molecule] = molecule
@@ -207,9 +208,11 @@ function broadcast_dict(dicts, key)
     return [dict[key] for dict in dicts]
 end
 
-function load_chains(init_path; args=Dict(), filename="", verbose=false)
+function load_chains(init_path; args=Dict(), filename="", verbose=false, fold=true)
     input_files = Vector{String}()
-    if isfile(init_path)
+    if init_path isa AbstractVector # to keep job ordered in order to well restart the simulations
+        append!(input_files,init_path)
+    elseif isfile(init_path)
         push!(input_files, init_path)
     elseif isdir(init_path)
         for (root, dirs, files) in walkdir(init_path)
@@ -281,7 +284,9 @@ function load_chains(init_path; args=Dict(), filename="", verbose=false)
     end
 
     # Fold back into the box
-    initial_position_array .= [[fold_back(x, box) for x in X] for (X, box) in zip(initial_position_array, initial_box_array)]
+    if fold
+        initial_position_array .= [[fold_back(x, box) for x in X] for (X, box) in zip(initial_position_array, initial_box_array)]
+    end
 
     # Copy configurations nsim times (replicas)
     if haskey(args, "nsim") && !isnothing(args["nsim"]) && args["nsim"] > 1

src/IO/exyz.jl

@@ -41,10 +41,12 @@ function read_header(data, format::EXYZ)
     end
     lattice_matrix = reshape(lattice_values, 3, 3)
     box = lattice_matrix[diagind(lattice_matrix)]
-    column_match = match(r"Properties=(.*)", metadata_line)
+    column_match = match(r"Properties=(\S+)", metadata_line)
     column_str = mat === nothing ? nothing : column_match.captures[1]
     column_info = parse_column_string(column_str, format)
-    return N, box, column_info, split(metadata_line, " ")
+    time_match = match(r"Time=(\d+)", metadata_line)
+    t = time_match === nothing ? 0 : parse(Int, time_match.captures[1])
+    return t, N, box, column_info, split(metadata_line, " ")
 end
 
 function get_selrow(::EXYZ, N, m)

src/IO/lammps.jl

@@ -77,12 +77,13 @@ function read_header(data, format::LAMMPS)
     number_of_atoms_index = findfirst(contains("ITEM: NUMBER OF ATOMS"), data)
     box_bounds_index = findfirst(contains("ITEM: BOX BOUNDS"), data)
     columns_index = findfirst(contains("ITEM: ATOMS"), data)
+    t = timestep_index === nothing ? 0 : parse(Int, data[timestep_index + 1]) # t value is the next line after "ITEM: TIMESTEP"
     N = parse(Int, data[number_of_atoms_index + 1])
     box_bounds = data[box_bounds_index + 1:box_bounds_index + 3]
     box_bounds = [parse.(Float64, split(elt)) for elt in box_bounds] # Convert row elements to Float64
     box = [elt[2] - elt[1] for elt in box_bounds]
     column_info = parse_column_string(data[columns_index], format)
-    return N, box, column_info, []
+    return t, N, box, column_info, [] # add t to the returned value
 end
 
 function write_header(io, system::Particles, t, format::LAMMPS, digits::Integer)

src/IO/xyz.jl

@@ -41,13 +41,15 @@ function read_header(data, format::XYZ)
     metadata = split(data[2], " ")  # Metadata split into an array
 
     # Extract cell vector from metadata
+    step_idx = findfirst(startswith("step:"), metadata)
+    t = step_idx === nothing ? 0 : parse(Int, replace(metadata[step_idx], "step:" => ""))
     cell_str = replace(metadata[findfirst(startswith("cell:"), metadata)], "cell:" => "")
     cell_vector = parse.(Float64, split(cell_str, ","))
     d = length(cell_vector)
     box = SVector{d}(cell_vector)
     column_str = replace(metadata[findfirst(startswith("columns:"), metadata)], "columns:" => "")
     column_info = parse_column_string(column_str, format; d=d)
-    return N, box, column_info, metadata
+    return t, N, box, column_info, metadata
 end
 
 function get_system_column(::Atoms, ::XYZ)

src/ParticlesMC.jl

@@ -9,6 +9,7 @@ module ParticlesMC
 using Arianna, StaticArrays, Transducers
 using Comonicon, TOML
 using Comonicon: @main
+using Serialization
 
 export Particles
 abstract type Particles <: AriannaSystem end
@@ -183,14 +184,39 @@ ParticlesMC implemented in Comonicon.
     burn = get(sim, "burn", 0)
     seed = sim["seed"]
     parallel = sim["parallel"]
+    wall_time = get(sim, "restart", Inf) # restart if specified in toml else Inf to run uncapped
     output_path = get(sim, "output_path", "./")
 
+    function restart_format(sim)
+        for output in get(sim,"output",[])
+            if output["algorithm"] == "StoreLastFrames"
+                return eval(Meta.parse("$(get(output,"fmt","XYZ"))()"))
+            end
+        end
+        return nothing
+    end
+
+    # detection of a restart or fresh start
+    restart_enabled = isfinite(wall_time)
+    fmt_ckpt  = restart_enabled ? restart_format(sim) : nothing
+    lastframe = isnothing(fmt_ckpt) ? "" : joinpath(output_path, "chains", "1", "lastframe$(fmt_ckpt.extension)")
+    t_start   = (fmt_ckpt !== nothing && isfile(lastframe)) ? load_configuration(lastframe)[:t] : 0
+    restart   = t_start > 0
+
     # Setup RNG and basic variables
 
     # optional field
 
+    if restart
+        chains_dir = joinpath(output_path, "chains")
+        cdirs = sort(readdir(chains_dir); by = s -> parse(Int, s))   # "1","2",…,"10" in NUMERIC order
+        load_path = [joinpath(chains_dir, c, "lastframe$(fmt_ckpt.extension)") for c in cdirs]
+    else
+        load_path = config
+    end
+
     if bonds !== nothing
-        chains = load_chains(config, args=Dict(
+        chains = load_chains(load_path, args=Dict(
             "temperature" => temperature,
             "density" => density,
             "model" => model,
@@ -199,16 +225,18 @@ ParticlesMC implemented in Comonicon.
             "bonds" => bonds,
         ),
         filename=filename,
+        fold=!restart,
         )
     else
-        chains = load_chains(config, args=Dict(
+        chains = load_chains(load_path, args=Dict(
             "temperature" => temperature,
             "density" => density,
             "model" => model,
             "list_type" => list_type,
             "list_parameters" => list_parameters,
         ),
         filename=filename,
+        fold=!restart,
         )
     end
     algorithm_list = []
@@ -281,7 +309,7 @@ ParticlesMC implemented in Comonicon.
             algorithm  = (
                 algorithm=ComputeRotation,
                 scheduler=sched,
-                theta_T=theta_T,
+                θ_T=theta_T,
             )
         else
             error("Unsupported observable algorithm: $alg")
@@ -324,6 +352,7 @@ ParticlesMC implemented in Comonicon.
                 path=output_path,
             )
         elseif alg == "PrintTimeSteps"
+
             algorithm = (
                 algorithm=eval(Meta.parse(alg)),
                 scheduler=sched,
@@ -335,10 +364,30 @@ ParticlesMC implemented in Comonicon.
     end
     M = 1
     path = joinpath(output_path)
-    simulation = Simulation(chains, algorithm_list, steps; path=path, verbose=true)
+    simulation = Simulation(chains, algorithm_list, steps; t_start=t_start, path=path, verbose=true)
+
+    if restart 
+        for c in eachindex(simulation.chains)
+            rng_path = joinpath(simulation.path,"chains",string(c),"rng_state.jls")
+            rng = open(rng_path,"r") do file
+                deserialize(file)
+            end
+            simulation.algorithms[1].rngs[c] = rng
+        end
+    end
 
     # Run the simulation
-    run!(simulation)
+    status = run!(simulation; wall_time=wall_time)
+
+    # Save RNG sequence to prevent bias on restart
+    for c in eachindex(simulation.chains)
+        rng_path = joinpath(simulation.path,"chains",string(c),"rng_state.jls")
+        open(rng_path,"w") do file 
+            serialize(file,simulation.algorithms[1].rngs[c])
+        end
+    end
+
+    exit(status == :need_restart ? 1 : 0) # if t did not reached steps then a 1 flag is exit in order to restart the simulation through a bash script
 
 end
 

src/rotation.jl

@@ -4,7 +4,7 @@
 # Two algorithms :
 # 1. ComputeRotation : updates system.Φ (simulation.observable)
 # 2. StorePhiTrajectory : writes system.Φ to disk
-# system.Φ[k][m] = rotation vector for molecule m under theta_T[k]
+# system.Φ[k][m] = rotation vector for molecule m under θ_T[k]
 
 using LinearAlgebra
 using StaticArrays
@@ -86,17 +86,17 @@ RotationState{T}() where {T} = RotationState{T}([], [], [], false)
 ##########                      ##########
 
 mutable struct ComputeRotation{T} <: AriannaAlgorithm
-    theta_T::Vector{T}                  # one per method
+    θ_T::Vector{T}                  # one per method
     states::Vector{RotationState{T}}    # one per chain
 end
 
 function ComputeRotation(chains;
-                     theta_T::Vector{Float64}=[π/4],
+                     θ_T::Vector{Float64}=[π/4],
                      scheduler::Vector{Int}=Int[],
                      kwargs...)
     n = length(chains) # number of chains
     states = [RotationState{Float64}() for _ in 1:n] # n rotation states independent
-    return ComputeRotation{Float64}(theta_T, states)
+    return ComputeRotation{Float64}(θ_T, states)
 end
 
 ##### Initialisation of the simulation #####
@@ -109,19 +109,29 @@ function Arianna.initialise(algorithm::ComputeRotation, simulation::Simulation)
         state  = algorithm.states[c]
         T      = typeof(system.temperature)
         N_mol  = system.Nmol
-        n_θ    = length(algorithm.theta_T)
+        n_θ    = length(algorithm.θ_T)
 
         state.R_bodyframe = Vector{SMatrix{3,3,T,9}}(undef, N_mol)  # allocate once
         get_all_body_frames!(state.R_bodyframe, system)              # fill state.R_bodyframe
 
-        state.R_ref       = [copy(state.R_bodyframe) for _ in 1:n_θ]
-        state.Φ_acc       = [[zero(SVector{3,T}) for _ in 1:N_mol] for _ in 1:n_θ]
-        state.initialized = true
-
-        resize!(system.Φ, n_θ)
-        for k in 1:n_θ
-            system.Φ[k] = [zero(SVector{3,T}) for _ in 1:N_mol]
+        if simulation.t_start > 0
+            _, _, _, _, R_ref, Φ_acc, Φ = read_phi_frame(joinpath(simulation.path,"chains", "$c", "lastphiframe.dat"))
+            state.R_ref = R_ref
+            state.Φ_acc = Φ_acc
+            resize!(system.Φ, n_θ)
+            for k in 1:n_θ
+                system.Φ[k] = Φ[k]
+            end
+
+        else 
+            state.R_ref       = [copy(state.R_bodyframe) for _ in 1:n_θ]
+            state.Φ_acc       = [[zero(SVector{3,T}) for _ in 1:N_mol] for _ in 1:n_θ]
+            resize!(system.Φ, n_θ)
+            for k in 1:n_θ
+                system.Φ[k] = [zero(SVector{3,T}) for _ in 1:N_mol]
+            end
         end
+        state.initialized = true
     end
 end
 
@@ -134,7 +144,7 @@ function Arianna.make_step!(simulation::Simulation, algorithm::ComputeRotation)
         state  = algorithm.states[c]
         N_mol  = system.Nmol
         get_all_body_frames!(state.R_bodyframe, system)   # no new vector created updates state.R_bodyframe
-        for (k, θ_T) in enumerate(algorithm.theta_T)
+        for (k, θ_T) in enumerate(algorithm.θ_T)
             for m in 1:N_mol
                 dR        = state.R_ref[k][m]' * state.R_bodyframe[m]
                 Φ_current = rotation_vector(dR)
@@ -151,6 +161,42 @@ end
 
 function Arianna.finalise(::ComputeRotation, ::Simulation) end
 
+##### Read a checkpoint ɸ frame #####
+##########                 ##########
+
+function read_phi_frame(path)
+    open(path) do file
+        t = parse(Int, split(readline(file), "=")[2])
+        N_mol = parse(Int, split(readline(file), "=")[2])
+        n_θ = parse(Int, split(readline(file), "=")[2])
+        θ_T = parse.(Float64, split(split(readline(file), "=")[2], ","))
+
+        R_ref = Vector{Vector{SMatrix{3,3,Float64,9}}}(undef, n_θ)
+        Φ_acc = Vector{Vector{SVector{3,Float64}}}(undef, n_θ)
+        Φ     = Vector{Vector{SVector{3,Float64}}}(undef, n_θ)
+
+        for k in 1:n_θ
+            R_ref[k] = read_block(file, N_mol, row_to_matrix)
+            Φ_acc[k] = read_block(file, N_mol, row_to_vector)
+            Φ[k] = read_block(file, N_mol, row_to_vector)
+        end
+    return (t, N_mol, n_θ, θ_T, R_ref, Φ_acc, Φ)
+    end
+end
+
+function read_block(file, N_mol, builder)
+    readline(file) # skip first # line
+    return [builder(split(readline(file), " ")) for _ in 1:N_mol]
+end
+
+function row_to_matrix(parts)
+    return SMatrix{3,3,Float64}(parse.(Float64, parts[2:10]))
+end
+
+function row_to_vector(parts)
+    return SVector{3,Float64}(parse.(Float64, parts[2:4]))
+end
+
 ##### Store rotation vector trajectory #####
 ##########                        ##########
 
@@ -188,12 +234,13 @@ function Arianna.initialise(algorithm::StorePhiTrajectories, simulation::Simulat
     simulation.verbose && println("Opening Φ trajectory files...")
     for c in eachindex(simulation.chains)
         system = simulation.chains[c]
+        writing_mode = simulation.t_start > 0 ? "a" : "w"
         n_θ    = length(system.Φ)
         algorithm.paths[c] = [joinpath(algorithm.dirs[c], "phitrajectories_$k.dat")
                                for k in 1:n_θ]
-        algorithm.files[c] = open.(algorithm.paths[c], "w")
+        algorithm.files[c] = open.(algorithm.paths[c], writing_mode)
     end
-    algorithm.store_first && Arianna.make_step!(simulation, algorithm)
+    (simulation.t_start == 0 && algorithm.store_first) && Arianna.make_step!(simulation, algorithm)
 end
 
 function Arianna.make_step!(simulation::Simulation, algorithm::StorePhiTrajectories)
@@ -218,6 +265,8 @@ end
 ##### Store last Φ frame for checkpoint/restart #####
 ##########                                    ##########
 
+
+
 struct StoreLastPhiFrame <: AriannaAlgorithm
     paths::Vector{String}   # one per chain: chains/c/rotation_checkpoint.dat
 
@@ -248,14 +297,14 @@ function Arianna.finalise(algorithm::StoreLastPhiFrame, simulation::Simulation)
         system = simulation.chains[c]
         state  = compute_rot.states[c]
         N_mol  = system.Nmol
-        n_θ    = length(compute_rot.theta_T)
+        n_θ    = length(compute_rot.θ_T)
 
         open(algorithm.paths[c], "w") do file
             # header
             println(file, "t=$(simulation.t)")
             println(file, "N_mol=$N_mol")
-            println(file, "n_theta=$n_θ")
-            println(file, "theta_T=$(join(compute_rot.theta_T, ','))")
+            println(file, "n_θ=$n_θ")
+            println(file, "θ_T=$(join(compute_rot.θ_T, ','))")
 
             # R_ref and Φ_acc
             for k in 1:n_θ