Ensemble simulations

[milankl]

@ningliu-iga With the following code you can run ensemble simulations in parallel, save as ensembles.jl and execute with

julia -p N ensembles.jl

where N is the number of processors you want to use.

import Distributed: @distributed

@everywhere using SpeedyWeather
@everywhere using Printf

n_samples = 10
spectral_grid = SpectralGrid(trunc=63,nlev=1,Grid=FullClenshawGrid)
orography = EarthOrography(spectral_grid,smoothing_truncation=340)
implicit = SpeedyWeather.ImplicitShallowWater(spectral_grid,α=0.5)
output = OutputWriter(spectral_grid,ShallowWater,output_dt=6,output_vars=[:u,:v,:pres])

@sync @distributed for sample in 1:n_samples
    output.id = @sprintf("%04d",sample)
    initial_conditions = SpeedyWeather.RandomWaves()
    model = ShallowWaterModel(;spectral_grid,orography,implicit,output,initial_conditions)
    simulation = initialize!(model)
    run!(simulation,n_days=10,output=true)
end

which prints

milan@dhcp-10-29-133-232 output % julia -p 2 ensembles.jl
Weather is speedy: run 0001 100%|█████████| Time: 0:00:07 (322.97 years/day)
Weather is speedy: run 0006 100%|█████████| Time: 0:00:07 (326.40 years/day)
Weather is speedy: run 0002 100%|█████████| Time: 0:00:05 (439.79 years/day)
Weather is speedy: run 0007 100%|█████████| Time: 0:00:05 (441.00 years/day)
Weather is speedy: run 0003 100%|█████████| Time: 0:00:05 (411.43 years/day)
Weather is speedy: run 0008 100%|█████████| Time: 0:00:05 (403.73 years/day)
Weather is speedy: run 0004 100%|█████████| Time: 0:00:05 (420.51 years/day)
Weather is speedy: run 0009 100%|█████████| Time: 0:00:05 (417.16 years/day)
Weather is speedy: run 0005 100%|█████████| Time: 0:00:05 (414.74 years/day)
Weather is speedy: run 0010 100%|█████████| Time: 0:00:05 (422.55 years/day)

So worker 1 took run 0001, worker 2 run 0006, then worker 1 again run 0002, etc. Every sample takes somewhat longer than the 5-7 seconds here as it excludes the initialization, which is currently redone for simplicity.

[milankl]

One of the problems above is that a given process actually reinitializes the whole model including all memory allocation. We're missing a fill!(::PrognosticVariables,0) function or similar. Maybe a general set!(::PrognosticVariables,...) like Oceananigans does it. Here's an example ensembles script that tries to avoid reallocation

import Distributed: @distributed

@everywhere using SpeedyWeather
@everywhere using Printf

n_samples = 50
spectral_grid = SpectralGrid(trunc=63,nlev=1,Grid=FullClenshawGrid)
orography = EarthOrography(spectral_grid,smoothing=false)
implicit = SpeedyWeather.ImplicitShallowWater(spectral_grid,α=0.5)
initial_conditions = SpeedyWeather.RandomWaves()

@everywhere initialised = false

@everywhere function reinitialize!(simulation::SpeedyWeather.Simulation)
    (;initial_conditions) = simulation.model
    (;prognostic_variables) = simulation
    SpeedyWeather.initial_conditions!(prognostic_variables,initial_conditions,simulation.model)
    return simulation
end

@sync @distributed for sample in 1:n_samples
    if ~initialised
        output = OutputWriter(spectral_grid,ShallowWater,output_dt=1/4,output_vars=[:u,:v,:pres,:orography])
        model = ShallowWaterModel(;spectral_grid,orography,output,initial_conditions,implicit);
        simulation = initialize!(model);
        println("Model initialised.")
    end

    output.id = @sprintf("%04d",sample)
    reinitialize!(simulation)
    run!(simulation,n_days=1/2,output=true)
end