Restart mechanism with a wall-time budget

.gitignore

@@ -17,4 +17,12 @@ data/*
 deps/build.log
 docs/build/*
 
+# Learning notes / restart docs — keep commits & PRs code-only.
+# (Already-tracked .md like README.md and docs/src/*.md stay tracked;
+#  this only prevents NEW .md from being committed. Use `git add -f` to override.)
+*.md
+
+# Generated test output and fixtures
+test/data/
+
 !config_0.xyz

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"
@@ -19,7 +20,7 @@ Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 
 [compat]
 Aqua = "0.8"
-Arianna = "0.2"
+Arianna = "0.2.2"
 Comonicon = "1.0"
 ComponentArrays = "0.15"
 ConcreteStructs = "0.2"
@@ -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