diff --git a/Project.toml b/Project.toml
index e13a077..c23f192 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "MEDYANSimRunner"
 uuid = "b58a3b99-22e3-44d1-b5ea-258f082a6fe8"
 authors = ["nhz2 <nhz2@cornell.edu>"]
-version = "0.5.2"
+version = "0.5.3"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
diff --git a/README.md b/README.md
index b5af1ab..ebae1f3 100644
--- a/README.md
+++ b/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]
diff --git a/src/run-sim.jl b/src/run-sim.jl
index 9836bd3..1b906c9 100644
--- a/src/run-sim.jl
+++ b/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
diff --git a/test/example/main.jl b/test/example/main.jl
index 20186f0..36dc2d5 100644
--- a/test/example/main.jl
+++ b/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
 
diff --git a/test/example/output-ref-1_10/a/traj/footer.json b/test/example/output-ref-1_10/a/traj/footer.json
index 0d4ee78..ed56d92 100644
--- a/test/example/output-ref-1_10/a/traj/footer.json
+++ b/test/example/output-ref-1_10/a/traj/footer.json
@@ -1,4 +1,4 @@
 {
     "steps": 11,
-    "prev_sha256": "22fa73eaa3944debf0d26bb9d9cebcecc5f3047a5ffed0c42e622003e562c71c"
+    "prev_sha256": "858a821e8b5c67b9c55e77064443bea55cf8be8d06fa745f21433e0664992966"
 }
\ No newline at end of file
diff --git a/test/example/output-ref-1_10/a/traj/snap1.zarr.zip b/test/example/output-ref-1_10/a/traj/snap1.zarr.zip
index 7a4e4e7..f0ba672 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap1.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap1.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap10.zarr.zip b/test/example/output-ref-1_10/a/traj/snap10.zarr.zip
index d782a24..26210c8 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap10.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap10.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap11.zarr.zip b/test/example/output-ref-1_10/a/traj/snap11.zarr.zip
index 45ecafc..69e229e 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap11.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap11.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap2.zarr.zip b/test/example/output-ref-1_10/a/traj/snap2.zarr.zip
index 8d1c0cc..8fefbac 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap2.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap2.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap3.zarr.zip b/test/example/output-ref-1_10/a/traj/snap3.zarr.zip
index 4d13c24..e8a206c 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap3.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap3.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap4.zarr.zip b/test/example/output-ref-1_10/a/traj/snap4.zarr.zip
index c9c9aad..481bb7c 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap4.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap4.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap5.zarr.zip b/test/example/output-ref-1_10/a/traj/snap5.zarr.zip
index 348261d..53d9afd 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap5.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap5.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap6.zarr.zip b/test/example/output-ref-1_10/a/traj/snap6.zarr.zip
index 5674d33..1b93d0f 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap6.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap6.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap7.zarr.zip b/test/example/output-ref-1_10/a/traj/snap7.zarr.zip
index 120a322..5ad91f5 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap7.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap7.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap8.zarr.zip b/test/example/output-ref-1_10/a/traj/snap8.zarr.zip
index a985148..c4e8573 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap8.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap8.zarr.zip differ
diff --git a/test/example/output-ref-1_10/a/traj/snap9.zarr.zip b/test/example/output-ref-1_10/a/traj/snap9.zarr.zip
index dfb936e..c4d4c52 100644
Binary files a/test/example/output-ref-1_10/a/traj/snap9.zarr.zip and b/test/example/output-ref-1_10/a/traj/snap9.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/footer.json b/test/example/output-ref-1_9/a/traj/footer.json
index 2460bfe..b9c6f9f 100644
--- a/test/example/output-ref-1_9/a/traj/footer.json
+++ b/test/example/output-ref-1_9/a/traj/footer.json
@@ -1,4 +1,4 @@
 {
     "steps": 11,
-    "prev_sha256": "127baaf98475d3f598ec5e78b1a6ab23462c214cfd60c7354e9b56260410b8b2"
+    "prev_sha256": "028031634b968ea546b390ae2fe13d1b2ebb36d03383f93998affe212d398809"
 }
\ No newline at end of file
diff --git a/test/example/output-ref-1_9/a/traj/snap1.zarr.zip b/test/example/output-ref-1_9/a/traj/snap1.zarr.zip
index de3ea9a..d37a58a 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap1.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap1.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap10.zarr.zip b/test/example/output-ref-1_9/a/traj/snap10.zarr.zip
index be86750..81441c4 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap10.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap10.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap11.zarr.zip b/test/example/output-ref-1_9/a/traj/snap11.zarr.zip
index 34b79c7..77ce92a 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap11.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap11.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap2.zarr.zip b/test/example/output-ref-1_9/a/traj/snap2.zarr.zip
index 8d6a2f7..1f3ec4d 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap2.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap2.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap3.zarr.zip b/test/example/output-ref-1_9/a/traj/snap3.zarr.zip
index f41c5d5..e81915e 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap3.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap3.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap4.zarr.zip b/test/example/output-ref-1_9/a/traj/snap4.zarr.zip
index 9914784..a215324 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap4.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap4.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap5.zarr.zip b/test/example/output-ref-1_9/a/traj/snap5.zarr.zip
index ddd3946..3533229 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap5.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap5.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap6.zarr.zip b/test/example/output-ref-1_9/a/traj/snap6.zarr.zip
index dd28fd8..dba52d6 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap6.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap6.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap7.zarr.zip b/test/example/output-ref-1_9/a/traj/snap7.zarr.zip
index 3432502..5bdf740 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap7.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap7.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap8.zarr.zip b/test/example/output-ref-1_9/a/traj/snap8.zarr.zip
index 975cbf2..e4e8209 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap8.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap8.zarr.zip differ
diff --git a/test/example/output-ref-1_9/a/traj/snap9.zarr.zip b/test/example/output-ref-1_9/a/traj/snap9.zarr.zip
index d10bfa4..a4abc62 100644
Binary files a/test/example/output-ref-1_9/a/traj/snap9.zarr.zip and b/test/example/output-ref-1_9/a/traj/snap9.zarr.zip differ