Open milankl opened 12 months ago
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
@ningliu-iga With the following code you can run ensemble simulations in parallel, save as
ensembles.jl
and execute withwhere
N
is the number of processors you want to use.which prints
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.