Functions & Types index

Core

Defines the fundamental data structure (SimPara) describing how simulation data is parameterized. Further it contains helpful utility functions (calc_n_coeff).

SpeedyWeatherEmulator.SimPara — Type

SimPara{F}

Container for core simulation parameters that define SpeedyWeather.jl data generation, resulting data shape and storage name.

Fields

trunc::Int: Spectral truncation of the barotropic model (e.g. 5 for T5).
n_data::Int: Number of stored data time steps after spin-up.
n_ic::Int: Number of simulated initial conditions (independent runs).
n_spinup::Int: Number of spin-up steps discarded before sampling.
t_step::Float32: Physical time step length.
initial_cond::F: Optional generator for initial conditions; if nothing, random ICs are used.
id_key::String: Additional identifier to disambiguate saved datasets with identical trunc, n_data and n_ic.

source

SpeedyWeatherEmulator.SimPara — Method

SimPara(;   trunc,
            n_data,
            n_ic,
            n_spinup = 9,
            t_step = 1.0,
            initial_cond = nothing,
            id_key = "")

Convenience constructor for SimPara.

Arguments

trunc::Int: Spectral truncation of the barotropic model (e.g. 5 for T5).
n_data::Int: Number of stored data time steps after spin-up.
n_ic::Int: Number of simulated initial conditions (independent runs).
n_spinup::Int = 9: Number of spin-up steps discarded before sampling (Default data sampling begins at t=10h).
t_step::Real = 1.0: Physical time step length.
initial_cond = nothing: Optional generator for initial conditions; if nothing, random ICs are used. Can also be a Function or any callable object.
id_key::String = "": Additional identifier to disambiguate saved datasets with identical (trunc, n_data, n_ic).

Returns

::SimPara: Container for simulation parameters that define the simulation and data storage.

Examples

sim_para = SimPara(trunc=5, n_data=50, n_ic=200, id_key="_test")

source

SpeedyWeatherEmulator.calc_n_coeff — Method

calc_n_coeff(; trunc::Int)

Calculate the number of complex spectral coefficients for a given spectral truncation.

Arguments

trunc::Int: Spectral truncation of the barotropic model (e.g. 5 for T5).

Returns

n_coeff::Int: Number of complex spectral coefficients (without splitting into real/imag parts).

source

Handles data generation and preparation. This includes creating raw simulation data with SpeedyWeather.jl (generate_raw_data), wrapping it into structured containers (SimData), and formatting it into train/validation/test sets (DataPairs, FormattedData).

SpeedyWeatherEmulator.generate_raw_data — Method

generate_raw_data(sim_para::SimPara; overwrite::Bool=false, path::String="")

Generate raw vorticity data with SpeedyWeather.jl based on the given simulation parameters.

Description

Runs the barotropic model with spectral truncation sim_para.trunc and time step sim_para.t_step.
For each initial condition (IC), creates a new run subfolder and stores the simulated vorticity.
If overwrite=false and data already exist, generation is canceled.

Arguments

sim_para::SimPara: Container for parameters that define the simulation and data storage.
overwrite::Bool = false: If true, delete existing data and regenerate. If false, aborts safely when data already exist.
path::String = "": Optional absolute path for data storage. If left empty, the function defaults to the package's internal data/<type> folder.

Returns

::Bool: Returns a boolean value that indicates if the data was generated successfully.

Notes

Spin-up steps (n_spinup) are run but not stored in later SimData.
If sim_para.initial_cond is not nothing, this function applies the given IC via set!(sim, vor=…).
Sometimes overwriting raw_data files is not possible because the folders are open/busy. Data Generation is then canceled.
Only prognostic variables are stored.

Examples

sim_para = SimPara(trunc=5, n_data=50, n_ic=200, id_key="123")
generate_raw_data(sim_para; overwrite=true)
# → creates data/raw_data/raw_data_T5_ndata50_IC200_ID123/run_0001/output.jld2, ..., run_0200/output.jld2

source

SpeedyWeatherEmulator.SimData — Type

SimData(sim_para::SimPara, path::String="")

Construct a SimData container by loading previously generated raw data and extracting the spectral vorticity coefficients time series into a consistent tensor layout.

This constructor:

infers (n_coeff, n_data, n_ic) from sim_para,
allocates the target array data::Array{Float32,3} with shape (2 * n_coeff, n_data, n_ic),
iterates over runs ic = 1:n_ic, loads output.jld2, and reads output_vector,
for each stored step step ∈ {n_spinup+1, …, n_spinup + n_data}:
- extracts the spectral vorticity vor,
- writes real(vor) to rows 1:n_coeff and imag(vor) to rows n_coeff+1:2n_coeff,
- stores at time index step - n_spinup.

Arguments

sim_para::SimPara{F}: Container for parameters that define the simulation and data storage; must match the generated raw data on disk.
path::String = "": Optional absolute path of stored raw_data. If left empty, the function defaults to the package's internal data/<type> folder.

Returns

::SimData{F, Array{Float32,3}}: Container holding simulation data and corresponding sim. parameters.

Preconditions

Expects raw data in data_path(sim_para; type="raw_data") with per-run subfolders run_0001, run_0002, … each containing output.jld2 with an output_vector.
Raw data should be created beforehand via generate_raw_data(sim_para; overwrite=false).

Notes

The leading factor 2 in the first dimension stacks real and imaginary parts.
The time indexing uses step - n_spinup to map stored steps to 1:n_data.

Examples

sim_para = SimPara(trunc=5, n_data=50, n_ic=200)
# after generate_raw_data(sim_para) has been called:
sim_data = SimData(sim_para)
# inspect shapes
size(sim_data.data)  # (2*n_coeff, n_data, n_ic)

source

SpeedyWeatherEmulator.SimData — Type

SimData{F, A<:AbstractArray{Float32, 3}}

Container for SpeedyWeather.jl vorticity data together with the simulation parameters.

Fields

sim_para::SimPara{F}: Container for parameters that define the simulation and data storage.
data::A: Vorticity tensor with shape (2 * n_coeff, n_data, n_ic) where n_coeff = calc_n_coeff(sim_para.trunc). The first n_coeff rows store the real parts, the next n_coeff rows the imaginary parts of the complex spectral coefficients.

Notes

The layout is column-major and optimized for contiguous slicing over time and ICs.
Dimensions are inferred from sim_para and remain consistent across the pipeline.

source

SpeedyWeatherEmulator.DataPairs — Type

DataPairs{A<:AbstractArray{Float32, 2}}

Container for paired data samples (x,y) = (vor(t), vor(t+Δt)), already split into training, validation and test sets.

Fields

x_train::A: Training inputs vor(t).
y_train::A: Training targets vor(t+Δt).
x_valid::A: Validation inputs vor(t).
y_valid::A: Validation targets vor(t+Δt).
x_test::A: Test inputs vor(t).
y_test::A: Test targets vor(t+Δt).

Notes

All matrices have the same row dimension = 2 * n_coeff.
Columns index over independent time-pairs and ICs.

source

SpeedyWeatherEmulator.FormattedData — Type

FormattedData{F, A<:AbstractArray{Float32, 2}}

Container for formatted data, i.e. paired vorticity samples (x,y) = (vor(t), vor(t+Δt)).

Fields

sim_para::SimPara{F}: Container for parameters that define the simulation and data storage.
data_pairs::DataPairs{A}: The split and paired data.

source

SpeedyWeatherEmulator.FormattedData — Method

FormattedData(  sim_data::SimData; 
                splits::NamedTuple{(:train, :valid, :test),<:Tuple{Vararg{Real,3}}} = 
                    (train=0.7, valid=0.15, test=0.15))

Construct FormattedData directly from SimData by pairing consecutive time steps and splitting them into train/validation/test sets.

Description

Builds (x,y) pairs as
- x = vor(t) = spectral vorticity state at time t,
- y = vor(t+Δt) at the next time step.
Reshapes all ICs and times into column vectors.
Splits the resulting pairs according to the fractions in splits.

Arguments

sim_data::SimData: Container holding simulation data and corresponding sim. parameters.
splits::NamedTuple{(:train, :valid, :test),<:Tuple{Vararg{Real,3}}}: Fractions for train-, valid- and test-set. Default = (0.7, 0.15, 0.15).

Returns

::FormattedData: Container holding formatted (paired) simulation data and corresponding sim. parameters.

Notes

The number of total pairs is (n_data - 1) * n_ic.
Splits are normalized so that train + valid + test = 1.

Examples

fd = FormattedData(sim_data; splits=(train=0.7, valid=0.2, test=0.1))
size(fd.data_pairs.x_train)  # (2*n_coeff, n_train)

source

Emulator

Handles emulator definition, normalization, training, and evaluation. This includes Z-score normalization utilities (ZscorePara, zscore, inv_zscore), core emulator types (NeuralNetwork, Emulator, Losses), training workflow (train_emulator), and evaluation against SpeedyWeather.jl reference data (compare_emulator).

SpeedyWeatherEmulator.ZscorePara — Type

ZscorePara{F<:AbstractVector{Float32}}

Container for the parameters of a Z-score transformation.

Fields

μ::F: Mean for each coefficient across samples.
σ::F: Std for each coefficient across samples.

Notes

Typically computed from the training set only to avoid data leakage.
For each coefficient indexed i: zi = (xi - μi) / (σi + eps)

source

SpeedyWeatherEmulator.inv_zscore — Method

inv_zscore(x::AbstractArray{Float32}, stats::ZscorePara{<:AbstractVector{Float32}})

Inverse Z-score transformation (restore original scale).

Arguments

x::AbstractArray{Float32}: Z-score normalized data.
stats::ZscorePara{<:AbstractVector{Float32}}: Parameters with mean μ and std σ.

Returns

::typeof(x): Data rescaled back to the original distribution.

Examples

stats = ZscorePara([0.0f0, 1.0f0], [1.0f0, 2.0f0])
z = Float32[0 -0.5; 1 0.5]
x = inv_zscore(z, stats)

source

SpeedyWeatherEmulator.zscore — Method

zscore(x::AbstractArray{Float32}, stats::ZscorePara{<:AbstractVector{Float32}})

Apply a Z-score transformation to data x using the parameters in stats.

Arguments

x::AbstractArray{Float32}: Input data (rows = coefficients, columns = samples).
stats::ZscorePara{<:AbstractVector{Float32}}: Parameters with mean μ and std σ.

Returns

::typeof(x): Z-score normalized data.

Notes

Each coefficient is transformed independently.
A small eps(Float32) is added to σ to avoid division by zero.

Examples

stats = ZscorePara([0.0f0, 1.0f0], [1.0f0, 2.0f0])
x = Float32[0 2; 1 3]
z = zscore(x, stats)

source

SpeedyWeatherEmulator.Emulator — Type

Emulator{A<:AbstractVector{Float32}}

Container for a trained (or in-progress) neural network emulator.

Fields

sim_para::SimPara: Simulation parameters of the dataset used for training, validation and testing.
chain::Flux.Chain: Neural network architecture and weights.
zscore_para::ZscorePara{A}: Normalization parameters (mean/std of training set).

source

SpeedyWeatherEmulator.Emulator — Method

(m::Emulator)(x::AbstractArray{Float32})

Convenience call overload. Apply the trained emulator to spectral coefficients at time t to predict coefficients at t + Δt.

Arguments

x::AbstractArray{Float32}: Spectral coefficients of vorticity at time t. First dimension must be 2 * n_coeff.

Returns

::typeof(x): Emulator prediction for a vorticity at t + Δt (same size as input array).

source

SpeedyWeatherEmulator.Emulator — Method

Emulator(sim_para::SimPara, nn::NeuralNetwork, zscore_para::ZscorePara)

Constructor for an Emulator. Builds a feed-forward network with ReLU activations according to the given NeuralNetwork specs.

Arguments

sim_para::SimPara: Simulation parameters used for generating the training data.
nn::NeuralNetwork: Parameters of the architecture (layer sizes).
zscore_para::ZscorePara: Normalization parameters of the training data.

Returns

::Emulator: A wrapped Flux model with normalization metadata.

source

SpeedyWeatherEmulator.Losses — Type

Losses{F}

Container for logging training and validation losses and training time.

Fields

sim_para::SimPara{F}: Simulation parameters of the dataset used.
train::Vector{Float32}: Training loss per batch.
valid::Vector{Float32}: Validation loss per batch.
bpe_train::Int: Batches per epoch (training set).
bpe_valid::Int: Batches per epoch (validation set).
training_time::Float32: Time needed for training the model in seconds.

source

SpeedyWeatherEmulator.Losses — Method

Losses(sim_para::SimPara, bpe_train, bpe_valid)

Constructor for an empty Losses container.

Arguments

sim_para::SimPara{F}: Simulation parameters of the dataset used.
bpe_train::Int: Batches per epoch (training set).
bpe_valid::Int: Batches per epoch (validation set).

Returns

::Losses{F}: Initialized container with empty loss vectors and training_time = 0.0.

source

SpeedyWeatherEmulator.NeuralNetwork — Type

NeuralNetwork

Container for the layer dimensions of a neural network.

Fields

io_dim::Int: Dimension of the input and output layer (e.g. number of spectral coefficients).
hidden_dim::Int: Dimension of each hidden layer.
n_hidden::Int: Number of hidden layers.

source

SpeedyWeatherEmulator.NeuralNetwork — Method

NeuralNetwork(; io_dim::Int=54, hidden_dim::Int=640, n_hidden::Int=1)

Convenience constructor for NeuralNetwork.

Arguments

io_dim::Int = 54: Dimension of the input and output layer (e.g. number of spectral coefficients). Standard is the T5 model.
hidden_dim::Int = 640: Dimension of each hidden layer. Optimal value, obtained by hyperparameter optimization.
n_hidden::Int = 1: Number of hidden layers. Optimal value, obtained by hyperparameter optimization.

Returns

::NeuralNetwork: Parameter container.

source

SpeedyWeatherEmulator.compare_emulator — Method

compare_emulator(   em::Emulator; 
                    x_test::AbstractArray{Float32, 2},
                    y_test::AbstractArray{Float32, 2},
                    n_it::Int=1,
                    output::Bool=false,
                    all_coeff::Bool=false,
                    id_em::Bool=false)

Compare emulator predictions against SpeedyWeather.jl reference data.

Description

Applies the emulator em to test inputs x_test and compares results to y_test for n_it timesteps.
Computes mean relative error per spectral coefficient: relerri = |ŷi - yi| / (|y_i| + ε) * 100.
Prints mean relative error (all coefficients averaged) and maximum mean relative error.
Optionally prints coefficient-wise relative errors.

Arguments

em::Emulator: Trained emulator to evaluate.
x_test::AbstractArray{Float32, 2}: Test inputs (vorticity coefficients at t) of form (2 * n_coeff, N).
y_test::AbstractArray{Float32, 2}: Reference outputs from SpeedyWeather.jl (at t+nit*Δt) of form (2 * ncoeff, N).
n_it::Int: Number of timesteps compared.
output::Bool=false: If true, print to STDOUT.
all_coeff::Bool=false: If true (and also output=true) print relative error for each coefficient.
id_em::Bool=false: The identity emulator is used (em(vor(t)) = vor(t))

Returns

mean_mean_rel::Float32: The mean (all spectral coeff.) mean (all possible datapairs) relative error for n_it timesteps.

Notes

The larger n_it, the fewer data pairs are available for comparison. For example: n_data=4 and n_it=2 leads to data pairs 1-2-3 and 2-3-4.
Some coefficients in SpeedyWeather.jl are structurally zero → flagged in output.
Errors are reported in percent [%].

Examples

em, losses = train_emulator(nn, fd)
compare_emulator(em; 
    x_test=fd.data_pairs.x_test,
    y_test=fd.data_pairs.y_test,
    all_coeff=true)

source

SpeedyWeatherEmulator.train_emulator — Method

train_emulator( nn::NeuralNetwork, 
                fd::FormattedData; 
                sim_para::SimPara=fd.sim_para, 
                batchsize::Int = 32, n_epochs::Int=300, η0::Real=0.001)

Train an emulator (neural network) with the given architecture and data.

Description

Computes Z-score parameters from the training set.
Constructs an Emulator with the given NeuralNetwork.
Applies Z-score normalization to training and validation pairs.
Trains the network using Adam with initial learning rate η0.
Halves the learning rate every 30 epochs.
Logs training and validation losses.
Calls compare_emulator on the test set after training.

Arguments

nn::NeuralNetwork: Defines the structure (layer sizes) of the neural network.
fd::FormattedData: Formatted dataset with train/valid/test splits of size (2*n_coeff, N).
sim_para::SimPara=fd.sim_para: Optional different simulation parameters for better saving.
batchsize::Int=32: Training batch size.
n_epochs::Int=300: Number of training epochs.
η0::Real=0.001: Initial learning rate.

Returns

em::Emulator: Trained emulator model (Flux chain + normalization).
losses::Losses: Recorded training/validation losses and batches per epoch.

Notes

Test set evaluation is not part of the training loop; only compare_emulator is called at the end.
Normalization statistics are always computed from the training set to avoid leakage.

Examples

nn = NeuralNetwork(io_dim=54, hidden_dim=128, n_hidden=2)
fd = FormattedData(sim_data; splits=(train=0.7, valid=0.2, test=0.1))
em, losses = train_emulator(nn, fd; batchsize=64, n_epochs=100, η0=0.0005)

source

Evaluation

Handles evaluation and visualization of emulator and simulation output. This includes plotting loss plots (plot_losses), and reconstructing vorticity fields from spectral coefficients as heatmaps (plot_heatmap).

SpeedyWeatherEmulator.plot_losses — Method

plot_losses(losses::Losses; title::String="Emulator Losses")

Plot training and validation losses stored in a Losses container.

Description

Plots training loss per batch (log-log scale).
Adds epoch-averaged training and validation loss curves.
Returns a Plots.Plot object for further customization or saving.

Arguments

losses::Losses: Container with training/validation loss history and number of batches per epoch.
title::String="Emulator Losses": Optional argument for different plot titles (e.g. different simulation parameters).

Returns

p::Plots.Plot: Combined plot of training and validation losses.

Notes

Training batches per epoch = losses.bpe_train.
Validation batches per epoch = losses.bpe_valid.
Number of epochs is inferred as length(losses.train) / bpe_train.

Examples

emu, losses = train_emulator(nn, fd)
p = plot_losses(losses)
display(p)

source

SpeedyWeatherEmulator.plot_heatmap — Method

plot_heatmap(vec::AbstractVector{Float32}; trunc::Int, title::String="", range::Tuple{Real,Real}=(-5e-5, +5e-5))

Plot a heatmap of a vorticity field reconstructed from a spectral coefficient vector.

Description

Converts spectral coefficients in vector form into a LowerTriangularMatrix.
Transforms this into a grid suitable for plotting.
Displays the grid as a heatmap using CairoMakie.

Arguments

vec::AbstractVector{Float32}: Spectral coefficient vector (real/imag stacked).
trunc::Int: Spectral truncation of the model (e.g. 5 for T5).
title::String = "": Optional argument for different plot titles (e.g. different simulation parameters).
range::Tuple{Real,Real} = (-5e-5, +5e-5): Defines the color range of the heatmap plot.

Returns

::CairoMakie.Heatmap: Heatmap figure object.

Examples

vec = rand(Float32, 54)   # random spectral coeffs for trunc=5
fig = plot_heatmap(vec; trunc=5, title="Vorticity field")
display(fig)

source

SpeedyWeatherEmulator.vec_to_ltm — Method

vec_to_ltm(vec::AbstractVector{Float32}, trunc::Int)

Convert a spectral coefficient vector into a LowerTriangularMatrix representation.

Description

Interprets the coefficient vector as complex spectral coefficients.
Places them into a triangular matrix layout consistent with SpeedyWeather.jl.

Arguments

vec::AbstractVector{Float32}: Vector of spectral coefficients.
trunc::Int: Spectral truncation of the model (e.g. 5 for T5).

Returns

L::LowerTriangularMatrix{ComplexF32}: Complex spectral coefficient matrix.

Notes

For trunc=5, produces an N=7 x M=6 LowerTriangularMatrix with 27 entries.
vec is expected to be structured as [Re(c1), …, Re(cncoeff), Im(c1), …, Im(cncoeff)].

Examples

n = calc_n_coeff(trunc=5)
vec = rand(Float32, 2*n)
L = vec_to_ltm(vec, 5)

source

Functions & Types index

Core

IO

Data

Emulator

Evaluation