loop output (#32)

* add optional output

* fix version

* typo

Project.toml

@@ -1,7 +1,7 @@
 name = "MEDYANSimRunner"
 uuid = "b58a3b99-22e3-44d1-b5ea-258f082a6fe8"
 authors = ["nhz2 <[email protected]>"]
-version = "0.5.2"
+version = "0.5.3"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

README.md

@@ -81,8 +81,11 @@ Also return the expected value of step when done will first be true, used for di
 
 This function should not mutate `state`
 
-#### `loop(step::Int, state; kwargs...) -> state`
-Return the state that gets passed to `save`
+#### `loop(step::Int, state; output::SmallZarrGroups.ZGroup, kwargs...) -> state`
+Return the state that gets passed to `save` and `load`
+
+Optionally, mutate the `output` keyword argument.
+When saving the snapshot, this group will get saved as `"out"`
 
 
 ### Main loop pseudo code
@@ -99,9 +102,10 @@ group = ZGroup(childern=Dict("snap" => save(step, state))
 SmallZarrGroups.save_zip(snapshot_zip_file, group)
 state = load(step, SmallZarrGroups.load_zip(snapshot_zip_file)["snap"], state)
 while true
-    state = loop(step, state)
+    output = ZGroup()
+    state = loop(step, state; output)
     step = step + 1
-    group = ZGroup(childern=Dict("snap" => save(step, state))
+    group = ZGroup(childern=Dict("snap"=>save(step, state), "out"=>output)
     SmallZarrGroups.save_zip(snapshot_zip_file, group)
     state = load(step, SmallZarrGroups.load_zip(snapshot_zip_file)["snap"], state)
     if done(step::Int, state)[1]

src/run-sim.jl

@@ -153,12 +153,13 @@ function start_job(out_dir, job::String;
             state, prev_sha256 = save_load_state!(rng_state, step, state, traj, save, load, prev_sha256)
             @info "Simulation started."
             while true
+                output = ZGroup()
                 copy!(Random.default_rng(), rng_state)
-                state = loop(step, state)
+                state = loop(step, state; output)
                 copy!(rng_state, Random.default_rng())
 
                 step += 1
-                state, prev_sha256 = save_load_state!(rng_state, step, state, traj, save, load, prev_sha256)
+                state, prev_sha256 = save_load_state!(rng_state, step, state, traj, save, load, prev_sha256, output)
 
                 copy!(Random.default_rng(), rng_state)
                 isdone::Bool, expected_final_step::Int64 = done(step::Int, state)
@@ -259,12 +260,13 @@ function continue_job(out_dir, job;
                 end
             end
             while true
+                output = ZGroup()
                 copy!(Random.default_rng(), rng_state)
-                state = loop(step, state)
+                state = loop(step, state; output)
                 copy!(rng_state, Random.default_rng())
 
                 step += 1
-                state, prev_sha256 = save_load_state!(rng_state, step, state, traj, save, load, prev_sha256)
+                state, prev_sha256 = save_load_state!(rng_state, step, state, traj, save, load, prev_sha256, output)
 
                 copy!(Random.default_rng(), rng_state)
                 isdone, expected_final_step = done(step::Int, state)
@@ -291,6 +293,7 @@ function save_load_state!(
         save,
         load,
         prev_sha256::String,
+        output::Union{Nothing, ZGroup}=nothing,
     )
     snapshot_group = ZGroup()
 
@@ -299,6 +302,9 @@ function save_load_state!(
     copy!(rng_state, Random.default_rng())
 
     snapshot_group["snap"] = sub_snapshot_group
+    if !isnothing(output)
+        snapshot_group["out"] = output
+    end
     attrs(snapshot_group)["rng_state"] = rng_2_str(rng_state)
     attrs(snapshot_group)["step"] = step
     attrs(snapshot_group)["prev_sha256"] = prev_sha256

test/example/main.jl

@@ -44,10 +44,11 @@ function done(step::Int, state; kwargs...)
     step > 10, 11
 end
 
-function loop(step::Int, state; kwargs...)
+function loop(step::Int, state; output::ZGroup, kwargs...)
     a = sum(state) + rand(0:1)
     state[1] = state[2]
     state[2] = a
+    attrs(output)["sum"] = sum(state)
     state
 end
 

test/example/output-ref-1_10/a/traj/footer.json

@@ -1,4 +1,4 @@
 {
     "steps": 11,
-    "prev_sha256": "22fa73eaa3944debf0d26bb9d9cebcecc5f3047a5ffed0c42e622003e562c71c"
+    "prev_sha256": "858a821e8b5c67b9c55e77064443bea55cf8be8d06fa745f21433e0664992966"
 }

test/example/output-ref-1_9/a/traj/footer.json

@@ -1,4 +1,4 @@
 {
     "steps": 11,
-    "prev_sha256": "127baaf98475d3f598ec5e78b1a6ab23462c214cfd60c7354e9b56260410b8b2"
+    "prev_sha256": "028031634b968ea546b390ae2fe13d1b2ebb36d03383f93998affe212d398809"
 }