New function offline_run! to write data during run!

CHANGELOG.md

@@ -1,5 +1,9 @@
 # main
 
+# v5.16
+
+- New function `offline_run!` allows writing data to file at predefined intervals during `run!` instead of storing it in memory. Currently supports [CSV](https://csv.juliadata.org/stable/) and [Arrow](https://apache.github.io/arrow-julia/stable/) files.
+
 # v5.15
 - Agents.jl moved to Julia 1.9+, and now exports visualization
   and interactive applications automatically once Makie (or Makie backends

Project.toml

@@ -1,11 +1,12 @@
 name = "Agents"
 uuid = "46ada45e-f475-11e8-01d0-f70cc89e6671"
 authors = ["George Datseris", "Tim DuBois", "Aayush Sabharwal", "Ali Vahdati", "Adriano Meligrana"]
-version = "5.15.2"
+version = "5.16.0"
 
 [deps]
-CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
+Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -29,12 +30,12 @@ Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 OSMMakie = "76b6901f-8821-46bb-9129-841bc9cfe677"
 
 [extensions]
-AgentsVisualizations = "Makie"
 AgentsOSMVisualizations = "OSMMakie"
+AgentsVisualizations = "Makie"
 
 [compat]
-CommonSolve = "0.2.4"
 CSV = "0.9.7, 0.10"
+CommonSolve = "0.2.4"
 DataFrames = "0.21, 0.22, 1"
 DataStructures = "0.18"
 Distributions = "0.25"

docs/src/api.md

@@ -328,5 +328,10 @@ AgentsIO.populate_from_csv!
 AgentsIO.dump_to_csv
 ```
 
+It is also possible to write data to file at predefined intervals while running your model, instead of storing it in memory:
+```@docs
+offline_run!
+```
+
 In case you require custom serialization for model properties, refer to the [Developer Docs](@ref)
 for details.

src/simulations/collect.jl

@@ -1,4 +1,5 @@
 export run!,
+    offline_run!,
     collect_agent_data!,
     collect_model_data!,
     init_agent_dataframe,
@@ -34,6 +35,8 @@ Run the model (step it with the input arguments propagated into [`step!`](@ref))
 data specified by the keywords, explained one by one below. Return the data as
 two `DataFrame`s, one for agent-level data and one for model-level data.
 
+See also [`offline_run!`](@ref) to write data to file while running the model.
+
 ## Data-deciding keywords
 * `adata::Vector` means "agent data to collect". If an entry is a `Symbol`, e.g. `:weight`,
   then the data for this entry is agent's field `weight`. If an entry is a `Function`, e.g.
@@ -101,27 +104,22 @@ If `a1.weight` but `a2` (type: Agent2) has no `weight`, use
   or [`deepcopy`](https://docs.julialang.org/en/v1/base/base/#Base.deepcopy) if some data are
   nested mutable containers. Both of these options have performance penalties.
 * `agents_first=true` : Whether to update agents first and then the model, or vice versa.
-* `showprogress=false` : Whether to show a progress bar.
+* `showprogress=false` : Whether to show progress
 """
 function run! end
 
 run!(model::ABM, agent_step!, n::Int = 1; kwargs...) =
     run!(model::ABM, agent_step!, dummystep, n; kwargs...)
 
-function run!(
-    model,
-    agent_step!,
-    model_step!,
-    n;
-    when = true,
-    when_model = when,
-    mdata = nothing,
-    adata = nothing,
-    obtainer = identity,
-    agents_first = true,
-    showprogress = false,
-)
-
+function run!(model, agent_step!, model_step!, n;
+        when = true,
+        when_model = when,
+        mdata = nothing,
+        adata = nothing,
+        obtainer = identity,
+        agents_first = true,
+        showprogress = false,
+    )
     df_agent = init_agent_dataframe(model, adata)
     df_model = init_model_dataframe(model, mdata)
     if n isa Integer
@@ -164,6 +162,157 @@ function run!(
     return df_agent, df_model
 end
 
+"""
+    offline_run!(model, agent_step! [, model_step!], n::Integer; kwargs...)
+    offline_run!(model, agent_step!, model_step!, n::Function; kwargs...)
+
+Do the same as [`run`](@ref), but instead of collecting the whole run into an in-memory 
+dataframe, write the output to a file after collecting data `writing_interval` times and 
+empty the dataframe after each write.
+Useful when the amount of collected data is expected to exceed the memory available
+during execution.
+
+## Keywords
+* `backend=:csv` : backend to use for writing data.
+  Currently supported backends: `:csv`, `:arrow`
+* `adata_filename="adata.\$backend"` : a file to write agent data on.
+  Appends to the file if it already exists, otherwise creates the file.
+* `mdata_filename="mdata.\$backend"`: a file to write the model data on.
+  Appends to the file if it already exists, otherwise creates the file.
+* `writing_interval=1` : write to file every `writing_interval` times data collection
+  is triggered. If the `when` keyword is not set, this corresponds to writing to file
+  every `writing_interval` steps; otherwise, the data will be written every
+  `writing_interval` times the `when` condition is satisfied
+  (the same applies to `when_model`).
+"""
+function offline_run! end
+
+offline_run!(model::ABM, agent_step!, n::Int = 1; kwargs...) =
+    offline_run!(model::ABM, agent_step!, dummystep, n; kwargs...)
+
+function offline_run!(model, agent_step!, model_step!, n;
+        when = true,
+        when_model = when,
+        mdata = nothing,
+        adata = nothing,
+        obtainer = identity,
+        agents_first = true,
+        showprogress = false,
+        backend::Symbol = :csv,
+        adata_filename = "adata.$backend",
+        mdata_filename = "mdata.$backend",
+        writing_interval = 1,
+    )
+    df_agent = init_agent_dataframe(model, adata)
+    df_model = init_model_dataframe(model, mdata)
+    if n isa Integer
+        if when == true
+            for c in eachcol(df_agent)
+                sizehint!(c, n)
+            end
+        end
+        if when_model == true
+            for c in eachcol(df_model)
+                sizehint!(c, n)
+            end
+        end
+    end
+
+    run_and_write!(model, agent_step!, model_step!, df_agent, df_model, n;
+        when, when_model,
+        mdata, adata,
+        obtainer, agents_first,
+        showprogress,
+        backend, adata_filename, mdata_filename, writing_interval
+    )
+end
+
+function run_and_write!(model, agent_step!, model_step!, df_agent, df_model, n;
+    when, when_model,
+    mdata, adata,
+    obtainer, agents_first,
+    showprogress,
+    backend, adata_filename, mdata_filename, writing_interval
+)
+    s = 0
+    p = if typeof(n) <: Int
+        ProgressMeter.Progress(n; enabled=showprogress, desc="run! progress: ")
+    else
+        ProgressMeter.ProgressUnknown(desc="run! steps done: ", enabled=showprogress)
+    end
+
+    writer = get_writer(backend)
+    agent_count_collections = 0
+    model_count_collections = 0
+    while until(s, n, model)
+        if should_we_collect(s, model, when)
+            collect_agent_data!(df_agent, model, adata, s; obtainer)
+            agent_count_collections += 1
+            if agent_count_collections % writing_interval == 0
+                writer(adata_filename, df_agent, isfile(adata_filename))
+                empty!(df_agent)
+            end
+        end
+        if should_we_collect(s, model, when_model)
+            collect_model_data!(df_model, model, mdata, s; obtainer)
+            model_count_collections += 1
+            if model_count_collections % writing_interval == 0
+                writer(mdata_filename, df_model, isfile(mdata_filename))
+                empty!(df_model)
+            end
+        end
+        step!(model, agent_step!, model_step!, 1, agents_first)
+        s += 1
+        ProgressMeter.next!(p)
+    end
+
+    if should_we_collect(s, model, when)
+        collect_agent_data!(df_agent, model, adata, s; obtainer)
+        agent_count_collections += 1
+    end
+    if should_we_collect(s, model, when_model)
+        collect_model_data!(df_model, model, mdata, s; obtainer)
+        model_count_collections += 1
+    end
+    # catch collected data that was not yet written to disk
+    if !isempty(df_agent)
+        writer(adata_filename, df_agent, isfile(adata_filename))
+        empty!(df_agent)
+    end
+    if !isempty(df_model)
+        writer(mdata_filename, df_model, isfile(mdata_filename))
+        empty!(df_model)
+    end
+
+    ProgressMeter.finish!(p)
+    return nothing
+end
+
+"""
+    get_writer(backend)
+Return a function to write to file using a given `backend`.
+The returned writer function will take three arguments:
+filename, data to write, whether to append to existing file or not.
+"""
+function get_writer(backend)
+    @assert backend in (:csv, :arrow) "Backend $backend not supported."
+    if backend == :csv
+        return writer_csv
+    elseif backend == :arrow
+        return writer_arrow
+    end
+end
+
+writer_csv(filename, data, append) = AgentsIO.CSV.write(filename, data; append)
+
+function writer_arrow(filename, data, append)
+    if append
+        AgentsIO.Arrow.append(filename, data)
+    else
+        AgentsIO.Arrow.write(filename, data; file = false)
+    end
+end
+
 ###################################################
 # core data collection functions per step
 ###################################################

src/submodules/io/AgentsIO.jl

@@ -11,4 +11,6 @@ using Agents
 
 include("csv_integration.jl")
 include("jld2_integration.jl")
+include("arrow_integration.jl")
+
 end

src/submodules/io/arrow_integration.jl

@@ -0,0 +1,14 @@
+using Arrow
+using DataFrames
+
+# TODO: Implement populate_from and dump_to functions for Arrow.jl
+
+function populate_from_arrow!()
+    @error "Not yet implemented."
+    return
+end
+
+function dump_to_arrow!()
+    @error "Not yet implemented."
+    return
+end

test/collect_tests.jl

@@ -202,6 +202,65 @@ end
             mdata = [(m) -> (m.deep.data[i]) for i in 1:length(model.deep.data)],
         )
         @test Array{Float64,1}(model_data[1, 2:end]) == model.deep.data
+
+        @testset "Writing to file while running" begin
+
+            # CSV
+            offline_run!(model, agent_step!, model_step!, 365 * 5;
+                when_model = each_year,
+                when = six_months,
+                mdata = [:flag, :year],
+                adata = [(:weight, mean)],
+                writing_interval = 3
+            )
+
+            adata_saved = CSV.read("adata.csv", DataFrame)
+            @test size(adata_saved) == (11, 2)
+            @test propertynames(adata_saved) == [:step, :mean_weight]
+
+            mdata_saved = CSV.read("mdata.csv", DataFrame)
+            @test size(mdata_saved) == (6, 3)
+            @test propertynames(mdata_saved) == [:step, :flag, :year]
+
+            rm("adata.csv")
+            rm("mdata.csv")
+            @test !isfile("adata.csv")
+            @test !isfile("mdata.csv")
+
+            # Arrow
+            offline_run!(model, agent_step!, model_step!, 365 * 5;
+                when_model = each_year,
+                when = six_months,
+                backend = :arrow,
+                mdata = [:flag, :year],
+                adata = [(:weight, mean)],
+                writing_interval = 3
+            )
+
+            adata_saved = DataFrame(Arrow.Table("adata.arrow"))
+            @test size(adata_saved) == (11, 2)
+            @test propertynames(adata_saved) == [:step, :mean_weight]
+
+            mdata_saved = DataFrame(Arrow.Table("mdata.arrow"))
+            @test size(mdata_saved) == (6, 3)
+            @test propertynames(mdata_saved) == [:step, :flag, :year]
+
+            @test size(vcat(DataFrame.(Arrow.Stream("adata.arrow"))...)) == (11, 2)
+            @test size(vcat(DataFrame.(Arrow.Stream("mdata.arrow"))...)) == (6, 3)
+
+            rm("adata.arrow")
+            rm("mdata.arrow")
+            @test !isfile("adata.arrow")
+            @test !isfile("mdata.arrow")
+
+            # Backends
+            @test_throws TypeError begin
+                offline_run!(model, agent_step!, model_step!, 365 * 5; backend = "hdf5")
+            end
+            @test_throws AssertionError begin
+                offline_run!(model, agent_step!, model_step!, 365 * 5; backend = :hdf5)
+            end
+        end
     end
 
     @testset "Low-level API for Collections" begin

test/runtests.jl

@@ -1,12 +1,13 @@
-using Test, Agents, Random, LinearAlgebra
+using Test, Agents, Random, LinearAlgebra, CSV
+using Agents.AgentsIO.Arrow
 using Agents.Graphs, Agents.DataFrames
 using StatsBase: mean
 using StableRNGs
 
 using Distributed
 addprocs(2)
 @everywhere begin
-    using Test, Agents, Random
+    using Test, Agents, Random, LinearAlgebra, CSV
     using Agents.Graphs, Agents.DataFrames
     using StatsBase: mean
     using StableRNGs