Compatibility with DrWatson.jl

[MilesCranmer]

Hey all,

Awesome software, I am really enjoying using it.

I was wondering if you have any recommendations for compatibility with DrWatson.jl? In particular, I am hoping to combine the tagsave functionality (which stores meta information about a simulation, like git commit, filename, etc) with JLD2OutputWriter somehow. This is what I have tried so far. (I'll give the full script at the end).

function init_jld2!(file, model)
    tag!(file)  # Tag with git information
    merge!(file, d)  # Save parameters to model file
end

sim.output_writers[:fields] = JLD2OutputWriter(
    model,
    (; ω, speed);
    # Save to `/data/simulations` based on parameters:
    filename=datadir("simulations", savename(d, "jld2")),
    schedule=TimeInterval(Δt * 3),
    init=init_jld2!
)

However, I get an error from tag!:

┌ Warning: Initialization of /Users/mcranmer/Documents/oceananigans_interpret/data/simulations/initial_condition=random_seed=0_stop_time=10.0_timestepper=RungeKutta3_Δt=0.2_Δx=0.00781.jld2 failed because MethodError: MethodError: no method matching tag!(::JLD2.JLDFile{JLD2.MmapIO})
│ Closest candidates are:
│   tag!(!Matched::AbstractDict{K, T}; gitpath, force, source, storepatch) where {K, T} at ~/.julia/packages/DrWatson/I8SbQ/src/saving_tools.jl:200

Do you have any recommendations for using DrWatson.jl with Oceananigans.JL? Maybe there's something basic I'm missing here about saving data files.

Thanks!
Miles

Full script

using DrWatson
@quickactivate "Learning Oceananigans.jl"

using Oceananigans
using Statistics
using Random


function makesim(d::Dict)
    @unpack size,
    extent,
    topology,
    architecture,
    timestepper,
    advection,
    closure,
    Δt,
    stop_time,
    Δx,
    initial_condition,
    seed = d

    actual_size = (size..., 1)

    grid = RectilinearGrid(; size, extent, topology)
    model = NonhydrostaticModel(; grid, timestepper, advection, closure)

    # Initial conditions:
    rng = MersenneTwister(seed)
    u, v = if initial_condition == "random"
        u = rand(rng, actual_size...)
        v = rand(rng, actual_size...)
        u .-= mean(u)
        v .-= mean(v)
        u, v
    else
        error("Unknown initial condition: $initial_condition")
    end

    set!(model; u=u, v=v)

    sim = Simulation(model; Δt, stop_time)
    sim.callbacks[:progress] = Callback(
        s -> @info "Iteration $(iteration(s)) at time $(time(s))", IterationInterval(100)
    )

    ω, speed = begin
        u, v, _ = model.velocities
        ω = ∂x(v) - ∂y(u)
        speed = sqrt(u^2 + v^2)
        ω, speed
    end

    function init_jld2!(file, model)
        tag!(file)  # Tag with git information
        merge!(file, d)  # Save parameters to model file
    end
    
    sim.output_writers[:fields] = JLD2OutputWriter(
        model,
        (; ω, speed);
        filename=datadir("simulations", savename(d, "jld2")),
        overwrite_existing=true,
        schedule=TimeInterval(Δt * 3),
        init=init_jld2!
    )
    return sim
end

Δt = 0.2
size = (128, 128)
extent = (1, 1)
Δx = extent[1] / size[1]
stop_time = 10.0
topology = (Periodic, Periodic, Flat)
architecture = CPU()
timestepper = :RungeKutta3
advection = nothing
closure = nothing
initial_condition = "random"
seed = 0

params = @strdict size extent topology architecture timestepper advection closure Δt stop_time Δx initial_condition seed

sim = makesim(params)
run!(sim)

[MilesCranmer]

Nevermind, I think I got it. The modification is as follows:

    function init_jld2!(file, model)
        output_dict = Dict{Symbol,Any}()
        tag!(output_dict)
        for (k, v) in output_dict
            file[string(k)] = v
        end
        for (k, v) in d
            file[string(k)] = v
        end
    end

I'll close this but please let me know if there's a better way to use these two packages together.

Cheers,
Miles

[glwagner]

Neat! Thanks for figuring that out. I'm going to convert this to a discussion because I think it would help your solution retain visibility for future users.

In the above I think the second loop over d is vestigial; you only need the loop over output_dict.

It may not be a terrible idea to do this by default...

[MilesCranmer]

d is the simulation parameters; I’m not sure you would want to add the git information to that with tag!, since it would appear in the filename and be passed to the model. What do you think? Or did you mean that Oceananigans already saves the simulation parameters to the file? In which case, the only reason for embedding them at the top level like this is that DrWatson can see them (and use them to search/filter experiments) - maybe there’s a way to have it recognize them though.

It may not be a terrible idea to do this by default

That would be nice! I’m new to both this package and DrWatson, but integration with it seems beneficial for running reproducible experiments.

[glwagner]

Though I realize I was a little confused --- it's not the Oceanangians version / git tag that's saved right? It's info about the user's personal repo.

[glwagner]

d is the simulation parameters

Do you mean things like time step, stop time, stop iteration, etc?

[MilesCranmer]

Correct and correct!

[MilesCranmer]

But a DrWatson.jl project commits your Manifest.toml file, so you already know the Oceananigans version anyways