Stability Issues with Multiple Weather Trajectories Generation

[CakeMixCore]

Context

I'm working on generating multiple weather trajectories using SpeedyWeather.jl for research on spectral decomposition of meshes and their PDEs. Currently experiencing stability issues where approximately 80% of runs fail with "NaN, or Inf detected" around 500-time steps or more, even with parameters close to the default Jablonowski-Williamson test case.

Current Setup

Using a modified version of the Jablonowski-Williamson test case with:

• T85 truncation
• FullClenshawGrid
• Dealiasing factor of 1
• 5 vertical layers
• 30-minute timestep

Parameters are varied within tight ranges:

• η₀: 0.251-0.253
• u₀: 34-35 m/s
• Perturbation latitude: 30°-50°
• Temperature variations kept minimal (ΔT: 0-4K, Tmin: 199-200K)

Questions for Discussion

1. Could the choice of FullClenshawGrid with dealiasing=1 be contributing to the instability? Would switching to a different grid type (e.g., HEALPixGrid) with higher dealiasing factors help? I tried with HEALPixGrid and dealising with 2, and still got failing runs.

2. Given that we're interested in high-frequency signals/turbulence up to a certain spherical harmonic degree, what would be the optimal combination of truncation level and dealiasing factor to achieve this while maintaining stability? We would like to guarantee, that these frequncies are present in the signal, so we can try to reconstruct it.

3. Would alternative initialization functions provide better stability while preserving or generating high-frequency components?

   • Current: ZonalWind + Jablonowski setup
   • Alternatives:
     • RossbyHaurwitzWave: Controlled wave patterns, stable evolution
     • ZonalJet: Physical jet stream dynamics, adjustable perturbations

Code Example

module WeatherSimulation

export run_simulation

# Core imports
using Random
using NetCDF
using JLD2
using SpeedyWeather
using Logging

# Model configuration
const MODEL_CONFIG = Dict(
    "trunc" => 85,
    "grid" => FullClenshawGrid,
    # "grid" => HEALPixGrid,
    "dealiasing" => 1,
    "nlayers" => 5,
    "timestep" => Minute(30)
)

# Default values for seed=1
const DEFAULT_VORDIV = Dict(
    "η₀" => 0.252,
    "u₀" => 35,
    "perturb_lat" => 40.0,
    "perturb_lon" => 20.0,
    "perturb_uₚ" => 1.0,
    "perturb_radius" => 1/10
)

const DEFAULT_TEMP = Dict(
    "η₀" => 0.252,
    "σ_tropopause" => 0.2,
    "u₀" => 35,
    "ΔT" => 0,
    "Tmin" => 200
)

function setup_model()
    spectral_grid = SpectralGrid(
        trunc=MODEL_CONFIG["trunc"],
        Grid=MODEL_CONFIG["grid"],
        dealiasing=MODEL_CONFIG["dealiasing"],
        nlayers=MODEL_CONFIG["nlayers"]
    )
    return Leapfrog(spectral_grid, Δt_at_T31=MODEL_CONFIG["timestep"])
end

function custom_zonal_wind(seed::Int;
    η₀ = nothing,
    u₀ = nothing,
    perturb_lat = nothing,
    perturb_lon = nothing,
    perturb_uₚ = nothing,
    perturb_radius = nothing
)
    use_defaults = isone(seed)
    return ZonalWind(
        η₀ = use_defaults ? DEFAULT_VORDIV["η₀"] : (isnothing(η₀) ? rand(0.251:0.001:0.253) : η₀),     # Keep within 0.001 of default
        u₀ = use_defaults ? DEFAULT_VORDIV["u₀"] : (isnothing(u₀) ? rand(34:35) : u₀),                 # Keep very close to default 35 m/s
        perturb_lat = use_defaults ? DEFAULT_VORDIV["perturb_lat"] : (isnothing(perturb_lat) ? rand(30:50) : perturb_lat),  # Mid-latitudes only
        perturb_lon = use_defaults ? DEFAULT_VORDIV["perturb_lon"] : (isnothing(perturb_lon) ? rand(0:359) : perturb_lon),  # Full range ok
        perturb_uₚ = use_defaults ? DEFAULT_VORDIV["perturb_uₚ"] : (isnothing(perturb_uₚ) ? rand(0.9:0.05:1.0) : perturb_uₚ),  # Keep close to 1.0
        perturb_radius = use_defaults ? DEFAULT_VORDIV["perturb_radius"] : (isnothing(perturb_radius) ? rand(0.09:0.01:0.1) : perturb_radius)  # Keep close to 0.1
    )
end

function custom_jablonowski_temperature(seed::Int;
    η₀ = nothing,
    σ_tropopause = nothing,
    u₀ = nothing,
    ΔT = nothing,
    Tmin = nothing
)
    use_defaults = isone(seed)
    base_u₀ = use_defaults ? DEFAULT_TEMP["u₀"] : (isnothing(u₀) ? rand(34:35) : u₀)  # Match zonal wind exactly
    return JablonowskiTemperature(
        η₀ = use_defaults ? DEFAULT_TEMP["η₀"] : (isnothing(η₀) ? rand(0.251:0.001:0.253) : η₀),      # Match vordiv η₀ exactly
        σ_tropopause = use_defaults ? DEFAULT_TEMP["σ_tropopause"] : (isnothing(σ_tropopause) ? rand(0.198:0.001:0.202) : σ_tropopause),  # Keep very close to 0.2
        u₀ = base_u₀,                # Exactly match vordiv u₀
        ΔT = use_defaults ? DEFAULT_TEMP["ΔT"] : (isnothing(ΔT) ? rand(0:2:4) : ΔT),                  # Keep very small
        Tmin = use_defaults ? DEFAULT_TEMP["Tmin"] : (isnothing(Tmin) ? rand(199:200) : Tmin)         # Keep very close to 200K
    )
end

function initialize_model(seed::Int)
    vordiv = custom_zonal_wind(seed)
    pres = PressureOnOrography()
    temp = custom_jablonowski_temperature(seed)
    return InitialConditions(;vordiv, pres, temp)
end

function save_initial_conditions(initial_conditions, output_file, seed)
    initial_condition_dict = Dict(
        "random_seed" => seed,
        "vordiv.eta" => initial_conditions.vordiv.η₀,
        "vordiv.u" => initial_conditions.vordiv.u₀,
        "vordiv.perturb_lat" => initial_conditions.vordiv.perturb_lat,
        "vordiv.perturb_lon" => initial_conditions.vordiv.perturb_lon,
        "vordiv.perturb_up" => initial_conditions.vordiv.perturb_uₚ,
        "vordiv.perturb_radius" => initial_conditions.vordiv.perturb_radius,
        # "pres.Deltap" => initial_conditions.pres.Δp,
        "temp.eta" => initial_conditions.temp.η₀,
        "temp.sigma_tropopause" => initial_conditions.temp.σ_tropopause,
        "temp.u" => initial_conditions.temp.u₀,
        "temp.DeltaT" => initial_conditions.temp.ΔT,
        "temp.Tmin" => initial_conditions.temp.Tmin
    )

    open(output_file, "w") do io
        for (key, value) in initial_condition_dict
            println(io, "$key: $value")
        end
    end
end

"""
Main simulation function that runs the weather model
"""
function move_simulation_files(speedy_dir::String, target_dir::String)
    # Move all files from SpeedyWeather's run folder to target directory
    for entry in readdir(speedy_dir)
        src = joinpath(speedy_dir, entry)
        dst = joinpath(target_dir, entry)
        mv(src, dst, force=true)
    end
    rm(speedy_dir)  # Remove the now-empty SpeedyWeather run folder
end

function run_simulation(output_dir::String, seed::Integer)
    Random.seed!(seed)

    # Ensure absolute path and create directory if it doesn't exist
    output_dir = abspath(output_dir)
    mkpath(output_dir)

    # Save initial conditions before changing directory
    initial_conditions = initialize_model(seed)
    save_initial_conditions(initial_conditions, 
        joinpath(output_dir, "initial_conditions.txt"), 
        seed)

    # Change to output directory for the simulation
    old_dir = pwd()
    cd(output_dir)

    try
        # Create and setup model
        spectral_grid = SpectralGrid(
            trunc=MODEL_CONFIG["trunc"],
            Grid=MODEL_CONFIG["grid"],
            dealiasing=MODEL_CONFIG["dealiasing"],
            nlayers=MODEL_CONFIG["nlayers"]
        )
        time_stepping = Leapfrog(spectral_grid, Δt_at_T31=MODEL_CONFIG["timestep"])

        # Create and initialize model
        model = PrimitiveDryModel(;
            spectral_grid,
            time_stepping,
            initial_conditions
        )

        # Initialize and run simulation
        simulation = initialize!(model, time=DateTime(1999,5,1))
        simulation.model.output.output_dt = Second(1800)
        simulation.model.output.run_path = output_dir

        # Run simulation
        run!(simulation, period=Day(24), output=true)

        # Move files if simulation completed
        for entry in readdir(output_dir)
            if startswith(entry, "run_") && isdir(joinpath(output_dir, entry))
                move_simulation_files(joinpath(output_dir, entry), output_dir)
                break
            end
        end
    finally
        cd(old_dir)
    end
end

end # module

Additional Information Running SpeedyWeather version: v0.12.0 Julia version: 1.11.1 Operating System: Linux (x86_64-linux-gnu)

Any guidance on achieving stable simulations while maintaining sufficient variability between trajectories would be highly appreciated.

[milankl]

So this runs for me without problems despite the dealiasing of 1

spectral_grid = SpectralGrid(trunc=85, nlayers=5, dealiasing=1)
model = PrimitiveWetModel(spectral_grid)
simulation = initialize!(model)
run!(simulation, period=Day(10))

# plot surface humidity
using CairoMakie
humid = simulation.diagnostic_variables.grid.humid_grid[:, 5];
heatmap(humid)

Note that this

time_stepping = Leapfrog(spectral_grid, Δt_at_T31=MODEL_CONFIG["timestep"])

        # Create and initialize model
        model = PrimitiveDryModel(;
            spectral_grid,
            time_stepping,
            initial_conditions
        )

with MODEL_CONFIG["timestep"] = Minute(30) is just the default anyway. It doesn't use a 30min time step (that would be too large at a resolution of T85 anyway) but it uses a timestep of 30min at T31 and scales it linearly to the resolution you chose. You can always check the actual time step if you just do model.time_stepping

julia> model.time_stepping
Leapfrog{Float32} <: SpeedyWeather.AbstractTimeStepper
├ trunc::Int64 = 85
├ nsteps::Int64 = 2
├ Δt_at_T31::Second = 1800 seconds
├ radius::Float32 = 6.371e6
├ adjust_with_output::Bool = true
├ robert_filter::Float32 = 0.1
├ williams_filter::Float32 = 0.53
├ Δt_millisec::Dates.Millisecond = 675000 milliseconds
├ Δt_sec::Float32 = 675.0
└ Δt::Float32 = 0.00010594883

so it's at T85 some ~11min because you at run with T85 at about 3x the resolution than T31.

[milankl]

Also perturbing the initial conditions

vordiv = ZonalWind(η₀=0.253, u₀=36, perturb_lat=45, perturb_lon=25, perturb_uₚ=2, perturb_radius=1/8)
temp = JablonowskiTemperature(η₀=0.253, u₀=36)
ic = InitialConditions(; vordiv, temp, humid = ConstantRelativeHumidity(), pres=PressureOnOrography())
model = PrimitiveWetModel(spectral_grid, initial_conditions=ic)
simulation = initialize!(model)
run!(simulation, period=Day(10))

as you can see 10 days aren't really enough to get a completely different weather ...

[milankl]

@CakeMixCore you put in a lot of code that I'm trying to understand what setup you're actually running. Could you provide a minimal example that blows up that would help me to reproduce the instability you're facing? Like

spectral_grid = SpectralGrid(?...
initial_conditions = InitialConditions(?...
other_components = ?
model = PrimitiveDryModel(all_of_the_above...)
simulation = intialize!(model)
run!(simulation, ...?)

you pull some random numbers but could you just pick some where the model blows up? btw, I doubt that the actual small parameter changes you do are crucial for stability but there might be something else going on that would be good to pin down

[CakeMixCore]

Thank you a lot for your help! With the PrimitiveWetModel I have no issues, and it seams like I can use a wider range of random initial conditions, which is perfect :)

The minimal example I provide explodes within the first few steps (tested it on two different operating systems):

# imports
using Random
using NetCDF
using JLD2
using SpeedyWeather

# Set up the spectral grid
spectral_grid = SpectralGrid(;
    trunc=85,
    Grid=FullClenshawGrid,
    nlayers=5,
    dealiasing=1
)

# Define initial conditions
vordiv = ZonalWind(
    η₀=0.24,
    u₀=32.0,
    perturb_lat=24.0,
    perturb_lon=292.0,
    perturb_uₚ=0.85,
    perturb_radius=0.1
)

temp = JablonowskiTemperature(
    η₀=0.26,
    σ_tropopause=0.19,
    u₀=38.0,
    ΔT=0.0,
    Tmin=200.0
)

initial_conditions = InitialConditions(
    vordiv=vordiv,
    temp=temp
)

# Time stepping configuration
time_stepping = Leapfrog(spectral_grid, Δt_at_T31=Minute(30))

# Create and initialize the model
model = PrimitiveDryModel(;
    spectral_grid,
    time_stepping,
    initial_conditions
)

# Initialize simulation
simulation = initialize!(model, time=DateTime(1999,5,1))
simulation.model.output.output_dt = Second(1800)

# Run simulation
run!(simulation, period=Day(24), output=true)

[milankl]

Maybe there's something off with the PrimitiveDryModel, e.g. for your setup something is unstable, so that would be good to investigate, already now:

1) I wouldn't change this

    η₀=0.24

as this "parameter" corresponds to a conversion in the vertical coordinates, as in, it's not supposed to be changed I believe. I have never tried changing it, so not sure whether that causes any problems.

2) if you don't provide initial conditions for a variable then it's zero

initial_conditions = InitialConditions(
    vordiv=vordiv,
    temp=temp,
    pres=PressureOnOrography()    # add this
)

because in the primitive models pres is actually the logarithm of surface pressure, starting with 0 is a ... bad idea. Maybe this is stupid and we should change the defaults here? Maybe ReferencePressure() would be good to implement that's in the case of the shallow water model actually ZeroInitially() but in the case of the primitive models actually the reference pressure from model.atmosphere.

[milankl]

I can reproduce the instability with your linear T85 FullClenshawGrid, but if you just use the default OctahedralGaussianGrid then it's fine. I haven't understood your argument about using as coarse a grid

We want to have as much high frequency signals/turbulences as possible, up to a certain spherical harmonic, so we can guarantee that there are no higher frequency components in a signal.

So in the end you want to use data on the grid but all of that data should contain only signals up to the truncation? This is the case in a spectral model regardless of what dealiasing you use. A time step of a simulation starts in spectral space where all the "memory" of the prognostic variables resides. Then transforms into grid space, so all information in grid space is only coming from spectral space. Then some calculations are done in grid space, where especially the non-linear terms can create higher wavenumbers those are then transformed back into spectral space, truncating back to the highest resolved wavenumber, then the timestepping is happening in spectral space before we start from the beginning. A higher dealiasing is crucial for the backtransform so that higher frequency signals of some diagnostic variables (!, not the models memory) are correctly projected back, but as long as you don't touch those variables I don't see your point in using a dealiasing of 1. It just makes the model less stable.

[milankl]

If you insist on using dealiasing = 1 then at least use a Gaussian grid (full or octahedral) for which the quadrature is more accurate (although not exact for deliasing<2) than the inexact clenshaw or healpix grids. Just tested, that also doesn't blow up

[CakeMixCore]

I totaly overlooked η₀, for sure I should not change it! Still having issues with the simulation, even fixed at 0.24 I fortgot to add the init for PressureOnOrography in the example code, but we normal used it.

We had initialy some troubles with other Grids in our downstream pipelines, but I think now other Grids would work as well, thank you for the hint!

Ahh, I see. Thank you, very insightful! I think we will then go for a higher dealiasing factor. I was not aware that in the simulation the dealiasing makes it much more stable.

Thanks a lot for your time, really appreciate it :)

[milankl]

Maybe answer the question for me what data of SpeedyWeather you'd like to use. Because in the end running a simulation and grabbing output can be on different grids/resolution. Also I get the impression that you just want to create an ensemble by perturbing the initial conditions, if that's the case I'd just add some small random vorticity to the initial conditions instead of messing the Jablonowski initial conditions otherwise. For example you could do

set!(simulation, vor=(λ, φ, σ) -> 1e-6*randn(), add=true)

after model initialization, maybe you can make the amplitude a little larger, play around with it?

[CakeMixCore]

At the end we use temp and vor, but considering the PrimitiveWetModel, we will also save cloud_top, humid and pres for each timestep and if available per level. We try to go for a resolution, that is between 15.000-35.000 mesh points.
At last we want to have different trajectories, but also want to save the initial conditions used in an txt file. When I use !set, can I get the used conditions from the prognostic/diagnostic variables?

[milankl]

cloud top is the height of the uppermost layer with clouds so that's 2D only. Pressure pres is the logarithm of surface pressure internally, but converted to hPa for output, so also only 2D, but you can multiply it by the sigma levels to get pressure on every vertical layer.

Before you hit run! you can store the initial conditions from simulation.prognostic_variables.vor for example but that's in spectral space (so lower triangular matrices of complex coefficients). The netcdf files will contain as a first time step those initial conditions in grid space. If you need these just to restart the model, do a run!(simulation, period=Day(0), output=true) before you simulate what you actually want to simulate as this will store a jld2 restart file that that you can reuse with initial_conditions=StartFromFile(id=????) (have a look into the documentation)