Skip to content

Cartesian spectra in ci #4053

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Cartesian spectra in ci #4053

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6fd6b73
rft: formatting
haakon-e Oct 11, 2025
9b35e1a
cartesian spectra
haakon-e Oct 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 90 additions & 51 deletions post_processing/ci_plots.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -348,8 +348,7 @@ function compute_spectrum(var::ClimaAnalysis.OutputVar; mass_weight = nothing)

FT = eltype(var.data)

mass_weight =
isnothing(mass_weight) ? ones(FT, length(var.dims[dim3])) : mass_weight
mass_weight = isnothing(mass_weight) ? ones(FT, length(var.dims[dim3])) : mass_weight

# Number of spherical wave numbers, excluding the first and the last
# This number was reverse-engineered from ClimaCoreSpectra
Expand All @@ -358,18 +357,14 @@ function compute_spectrum(var::ClimaAnalysis.OutputVar; mass_weight = nothing)
mesh_info = nothing

if !isempty(times)
output_spectrum =
zeros((length(times), num_output, length(var.dims[dim3])))
output_spectrum = zeros((length(times), num_output, length(var.dims[dim3])))
dims = Dict(time => times)
dim_attributes = Dict(time => var.dim_attributes[time])
for index in 1:length(times)
spectrum_data, _wave_numbers, _spherical, mesh_info =
power_spectrum_2d(FT, var.data[index, :, :, :], mass_weight)
output_spectrum[index, :, :] .=
dropdims(sum(spectrum_data, dims = 1), dims = 1)[
(begin + 1):(end - 1),
:,
]
dropdims(sum(spectrum_data, dims = 1), dims = 1)[(begin + 1):(end - 1), :]
end
else
dims = Dict{String, Vector{FT}}()
Expand All @@ -378,10 +373,7 @@ function compute_spectrum(var::ClimaAnalysis.OutputVar; mass_weight = nothing)
spectrum_data, _wave_numbers, _spherical, mesh_info =
power_spectrum_2d(FT, var.data[:, :, :], mass_weight)
output_spectrum[:, :] .=
dropdims(sum(spectrum_data, dims = 1), dims = 1)[
(begin + 1):(end - 1),
:,
]
dropdims(sum(spectrum_data, dims = 1), dims = 1)[(begin + 1):(end - 1), :]
end

w_numbers = collect(1:1:(mesh_info.num_spherical - 1))
Expand All @@ -392,17 +384,39 @@ function compute_spectrum(var::ClimaAnalysis.OutputVar; mass_weight = nothing)
dim_attributes[dim3] = var.dim_attributes[dim3]

attributes = Dict(
"short_name" => "log_spectrum_" * var.attributes["short_name"],
"long_name" => "Spectrum of " * var.attributes["long_name"],
"short_name" => "log_spectrum_" * short_name(var),
"long_name" => "Spectrum of " * long_name(var),
"units" => "",
)

return ClimaAnalysis.OutputVar(
attributes,
dims,
dim_attributes,
log10.(output_spectrum),
attributes, dims, dim_attributes, log10.(output_spectrum),
)
end

import ClimaCoreSpectra.FFTW
function compute_average_cartesian_spectrum_1d(x::ClimaAnalysis.OutputVar)
@assert length(x.dims) == 2 "Can only compute 1D or averaged 2D spectra"
@assert x.dim2index["time"] == 1 "Time must be first dimension"
space_name = x.index2dim[2]
nt, nspace = size(x.data)
Δ = only(unique(round.(Int, diff(x.dims[space_name])))) # Assume uniform grid, rounded to Int (m-scale)

x̂ = FFTW.rfft(x.data, 2) # FFT along non-time dim
power = mean(abs2.(x̂) ./ nspace, dims = 1) |> vec

# ω = FFTW.rfftfreq(nspace, 1/Δ)
ω = FFTW.rfftfreq(nspace, nspace) .+ 1 # wavenumber

attr = Dict(
"short_name" => "spectrum_" * short_name(x),
"long_name" => "Spectrum of " * long_name(x),
"units" => "",
)
dims = Dict("Log10 Wavenumber" => log10.(ω))
dim_attr = Dict("Log10 Wavenumber" => Dict("units" => "log(1/m)"))

return ClimaAnalysis.OutputVar(attr, dims, dim_attr, log10.(power))
end

"""
Expand Down Expand Up @@ -1287,13 +1301,15 @@ Helper function for `make_plots_generic`. Takes a list of variables and plots
them on the same axis.
"""
function plot_les_vert_profile!(grid_loc, var_group)
z = var_group[1].dims["z"]
units = var_group[1].attributes["units"]
var1 = var_group[1]
units = var1.attributes["units"]
z = var1.dims["z"]
z_units = var1.dim_attributes["z"]["units"]
ax = CairoMakie.Axis(
grid_loc[1, 1],
ylabel = "z [$(var_group[1].dim_attributes["z"]["units"])]",
xlabel = "$(short_name(var_group[1])) [$units]",
title = parse_var_attributes(var_group[1]),
ylabel = "z [$z_units]",
xlabel = "$(short_name(var1)) [$units]",
title = parse_var_attributes(var1),
)

for var in var_group
Expand All @@ -1315,48 +1331,71 @@ function make_plots(
"Dh_smag", "strainh_smag", # smag horizontal
"Dv_smag", "strainv_smag", # smag vertical
"edt", # DecayWithHeight vertical diffusivity
"husra", "hussn", # 1M microphysics
"cdnc", "ncra", # 2M microphysics
]
short_names = short_names ∩ collect(keys(simdirs[1].vars))
short_names_xyzt = short_names_xyzt ∩ collect(keys(simdirs[1].vars))

short_names_spectra = ["wa"]
short_names_spectra = short_names_spectra ∩ collect(keys(simdirs[1].vars))

# Window average from instantaneous snapshots?
function horizontal_average(var)
return average_xy(var)
end
function windowed_reduction(var)
hours = 3600.0
function average_t_last2hrs(var)
window_end = last(var.dims["time"])
window_start = window_end - 2hours
var_window = ClimaAnalysis.window(
var, "time"; left = window_start, right = window_end,
window_start = window_end - CA.time_to_seconds("2hours")
var_window = ClimaAnalysis.window(var, "time";
left = window_start, right = window_end,
)
var_reduced = horizontal_average(average_time(var_window))
return var_reduced
return average_xy(average_time(var_window))
end

var_groups_xyt_reduced =
map_comparison(simdirs, short_names) do simdir, short_name
return [get(simdir; short_name, reduction) |> windowed_reduction]
end
var_groups_z_avg_xyt = map_comparison(simdirs, short_names_xyzt) do simdir, short_name
return [get(simdir; short_name, reduction) |> average_t_last2hrs]
end

var_groups_xy_reduced =
map_comparison(simdirs, short_names) do simdir, short_name
return [get(simdir; short_name, reduction) |> horizontal_average]
end
var_groups_tz_avg_xy = map_comparison(simdirs, short_names_xyzt) do simdir, short_name
return [get(simdir; short_name, reduction) |> average_xy]
end

tmp_file = make_plots_generic(
output_paths,
var_groups_xyt_reduced,
output_name = "tmp";
tmp_file = make_plots_generic(output_paths, var_groups_z_avg_xyt;
output_name = "tmp",
plot_fn = plot_les_vert_profile!,
MAX_NUM_COLS = 2,
MAX_NUM_ROWS = 4,
MAX_NUM_COLS = 2, MAX_NUM_ROWS = 4,
)

make_plots_generic(
output_paths,
vcat(var_groups_xy_reduced...),
summary_file = make_plots_generic(output_paths, vcat(var_groups_tz_avg_xy...);
output_name = isempty(short_names_spectra) ? "summary" : "tmp2",
plot_fn = plot_parsed_attribute_title!,
summary_files = [tmp_file],
MAX_NUM_COLS = 2, MAX_NUM_ROWS = 4,
)

isempty(short_names_spectra) && return summary_file
# If spectra are available, make additional plots

var_groups_y_spectra =
map_comparison(simdirs, short_names_spectra) do simdir, short_name
var = get(simdir; short_name, reduction)
last_t = var.dims["time"][end]
var_ty = slice(var; x=50_000, z=10_000)
data_yt = window(var_ty, "time"; left = last_t - 10*3600)
compute_average_cartesian_spectrum_1d(data_yt)
end

var_groups_x_spectra =
map_comparison(simdirs, short_names_spectra) do simdir, short_name
var = get(simdir; short_name, reduction)
last_t = var.dims["time"][end]
var_tx = slice(var; y=50_000, z=10_000)
data_xt = window(var_tx, "time"; left = last_t - 10*3600)
compute_average_cartesian_spectrum_1d(data_xt)
end

make_plots_generic(
output_paths,
[var_groups_y_spectra; var_groups_x_spectra];
summary_files = [summary_file],
plot_fn = plot_spectrum_with_line!,
MAX_NUM_COLS = 2,
MAX_NUM_ROWS = 4,
)
Expand Down
Loading