Add some averaged diagnostics

Project.toml

@@ -75,7 +75,7 @@ Logging = "1"
 Makie = "0.24"
 NCDatasets = "0.14"
 NVTX = "1.0.3"
-Oceananigans = "0.104, 0.105, 0.106, 0.107"
+Oceananigans = "0.107"
 Poppler_jll = "24"
 Printf = "1"
 Random = "1"

config/longrun_configs/cmip_edonly_bucket.yml

@@ -19,6 +19,10 @@ mode_name: "cmip"
 nh_poly: 3
 nh_poly_coupler: 2
 ocean_model: "oceananigans"
+ocean_diagnostic_interval: "1days"
+ocean_diagnostic_mode: "averaged"
+seaice_diagnostic_interval: "1days"
+seaice_diagnostic_mode: "averaged"
 rmse_check: true
 share_surface_space: false
 start_date: "20100101"

experiments/CMIP/Manifest-v1.11.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.11.9"
 manifest_format = "2.0"
-project_hash = "bca350e948c0680799cbca66ccefa34cddffd425"
+project_hash = "5e0cfd7c746742c6cfe8842524f909d76456e222"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "bbc22a9a08a0ef6460041086d8a7b27940ed4ffd"
@@ -2398,9 +2398,9 @@ version = "0.5.5"
 
 [[deps.Oceananigans]]
 deps = ["Adapt", "Crayons", "CubedSphere", "Dates", "Distances", "DocStringExtensions", "FFTW", "GPUArraysCore", "Glob", "InteractiveUtils", "JLD2", "KernelAbstractions", "Krylov", "LinearAlgebra", "Logging", "MPI", "MuladdMacro", "OffsetArrays", "OrderedCollections", "Pkg", "Printf", "ReactantCore", "Rotations", "SeawaterPolynomials", "SparseArrays", "StaticArrays", "Statistics", "StructArrays"]
-git-tree-sha1 = "985c51ec987b5b479b88428473a6160f177ee83e"
+git-tree-sha1 = "b91e5bea8b0f2bf735f3c9b1a1a9ed64a22df8bf"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.107.5"
+version = "0.107.6"
 
     [deps.Oceananigans.extensions]
     OceananigansAMDGPUExt = ["AMDGPU", "AbstractFFTs"]

experiments/CMIP/Manifest.toml

@@ -1,8 +1,8 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.10.10"
+julia_version = "1.10.11"
 manifest_format = "2.0"
-project_hash = "86fefa5844d6025da2ae370d295a5ea10980deac"
+project_hash = "5bc9f7a9ad334cbb192b125e8dccececfd916bf6"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "bbc22a9a08a0ef6460041086d8a7b27940ed4ffd"
@@ -2170,7 +2170,7 @@ version = "1.1.10"
 [[deps.MbedTLS_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
-version = "2.28.2+1"
+version = "2.28.1010+0"
 
 [[deps.MeshArrays]]
 deps = ["CatViews", "Dates", "Distributed", "GeoInterface", "Glob", "LazyArtifacts", "NearestNeighbors", "Pkg", "Printf", "SharedArrays", "SparseArrays", "Statistics", "Unitful"]
@@ -2237,7 +2237,7 @@ version = "0.3.7"
 
 [[deps.MozillaCACerts_jll]]
 uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
-version = "2023.1.10"
+version = "2025.12.2"
 
 [[deps.MuladdMacro]]
 git-tree-sha1 = "cac9cc5499c25554cba55cd3c30543cff5ca4fab"
@@ -2385,9 +2385,9 @@ version = "0.5.5"
 
 [[deps.Oceananigans]]
 deps = ["Adapt", "Crayons", "CubedSphere", "Dates", "Distances", "DocStringExtensions", "FFTW", "GPUArraysCore", "Glob", "InteractiveUtils", "JLD2", "KernelAbstractions", "Krylov", "LinearAlgebra", "Logging", "MPI", "MuladdMacro", "OffsetArrays", "OrderedCollections", "Pkg", "Printf", "ReactantCore", "Rotations", "SeawaterPolynomials", "SparseArrays", "StaticArrays", "Statistics", "StructArrays"]
-git-tree-sha1 = "985c51ec987b5b479b88428473a6160f177ee83e"
+git-tree-sha1 = "b91e5bea8b0f2bf735f3c9b1a1a9ed64a22df8bf"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.107.5"
+version = "0.107.6"
 
     [deps.Oceananigans.extensions]
     OceananigansAMDGPUExt = ["AMDGPU", "AbstractFFTs"]
@@ -2450,7 +2450,7 @@ version = "0.3.29+0"
 [[deps.OpenBLAS_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
 uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
-version = "0.3.23+4"
+version = "0.3.23+5"
 
 [[deps.OpenEXR]]
 deps = ["Colors", "FileIO", "OpenEXR_jll"]
@@ -3834,7 +3834,7 @@ version = "2022.0.0+1"
 [[deps.p7zip_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
-version = "17.4.0+2"
+version = "17.6.1+0"
 
 [[deps.pixi_jll]]
 deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl"]

experiments/CMIP/Project.toml

@@ -65,7 +65,7 @@ MPI = "0.20"
 Makie = "0.24"
 NCDatasets = "0.14"
 NVTX = "1"
-Oceananigans = "0.104, 0.105, 0.106, 0.107"
+Oceananigans = "0.107"
 Poppler_jll = "24"
 SciMLBase = "~2.110, ~2.111, ~2.112"
 StaticArrays = "1.9.8"

ext/ClimaCouplerCMIPExt.jl

@@ -16,7 +16,13 @@ module ClimaCouplerCMIPExt
 
 import ClimaCoupler
 import ClimaCoupler:
-    Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Utilities, Plotting
+    Checkpointer,
+    FieldExchanger,
+    FluxCalculator,
+    Interfacer,
+    TimeManager,
+    Utilities,
+    Plotting
 import Oceananigans as OC
 import ClimaOcean as CO
 import ClimaSeaIce as CSI
@@ -34,4 +40,7 @@ include("ClimaCouplerCMIPExt/skin_temperature.jl")
 include("ClimaCouplerCMIPExt/oceananigans.jl")
 include("ClimaCouplerCMIPExt/clima_seaice.jl")
 
+include("ClimaCouplerCMIPExt/ocean_diagnostics.jl")
+include("ClimaCouplerCMIPExt/seaice_diagnostics.jl")
+
 end # module ClimaCouplerCMIPExt

ext/ClimaCouplerCMIPExt/clima_seaice.jl

@@ -75,6 +75,8 @@ function ClimaSeaIceSimulation(
     start_date = nothing,
     coupled_param_dict = CP.create_toml_dict(FT),
     dt = 5 * 60.0, # 5 minutes
+    seaice_diagnostic_interval = "1days",
+    seaice_diagnostic_mode = :averaged,
     extra_kwargs...,
 ) where {FT}
     # Initialize the sea ice with the same grid as the ocean
@@ -124,21 +126,6 @@ function ClimaSeaIceSimulation(
     # Since ocean and sea ice share the same grid, we can also share the remapping objects
     remapping = ocean.remapping
 
-    # Before version 0.96.22, the NetCDFWriter was broken on GPU
-    if arch isa OC.CPU
-        # Save all tracers and velocities to a NetCDF file at daily frequency
-        outputs = OC.prognostic_fields(ice.model)
-        jld_writer = OC.JLD2Writer(
-            ice.model,
-            outputs;
-            schedule = OC.TimeInterval(86400), # Daily output
-            filename = joinpath(output_dir, "seaice_diagnostics.jld2"),
-            overwrite_existing = true,
-            array_type = Array{FT},
-        )
-        ice.output_writers[:diagnostics] = jld_writer
-    end
-
     # Allocate space for the sea ice-ocean (io) fluxes
     io_bottom_heat_flux = OC.Field{OC.Center, OC.Center, Nothing}(grid)
     io_frazil_heat_flux = OC.Field{OC.Center, OC.Center, Nothing}(grid)
@@ -178,6 +165,13 @@ function ClimaSeaIceSimulation(
         model_Δt,
     )
 
+    add_seaice_diagnostics!(
+        sim;
+        output_dir,
+        interval = TimeManager.time_to_period(seaice_diagnostic_interval),
+        mode = seaice_diagnostic_mode,
+    )
+
     # Ensure ocean temperature is above freezing where there is sea ice
     CO.OceanSeaIceModels.above_freezing_ocean_temperature!(ocean.ocean, grid, ice)
     return sim

ext/ClimaCouplerCMIPExt/ocean_diagnostics.jl

@@ -0,0 +1,135 @@
+import Dates
+
+"""
+    diagnostic_schedule(mode, interval)
+
+Return an Oceananigans output schedule corresponding to `mode`:
+- `:averaged` → `OC.AveragedTimeInterval(interval)`
+- `:instantaneous` → `OC.TimeInterval(interval)`
+
+Throws for any other value.
+"""
+function diagnostic_schedule(mode::Symbol, interval)
+    if mode === :averaged
+        return OC.AveragedTimeInterval(interval)
+    elseif mode === :instantaneous
+        return OC.TimeInterval(interval)
+    else
+        error(
+            "Unknown diagnostic mode `$(mode)`. Expected `:averaged` or `:instantaneous`.",
+        )
+    end
+end
+
+"""
+    add_ocean_diagnostics!(ocean_sim::OceananigansSimulation;
+                           output_dir = joinpath("output_active", "clima_ocean"),
+                           interval = Dates.Day(1),
+                           mode = :averaged,
+                           filename_prefix = "ocean",
+                           file_splitting_interval = Dates.Day(15))
+
+Attach output writers to the underlying Oceananigans simulation inside an `OceananigansSimulation`.
+Two writers are added to `ocean_sim.ocean.output_writers`, both with the same `interval` and `mode`:
+
+1. **Surface diagnostics** (`<prefix>_surface.jld2`): 2-D fields —
+   SST, SSS, SSH, surface velocities, squared surface fields for variance, surface momentum/heat/freshwater fluxes.
+2. **3-D field diagnostics** (`<prefix>_fields.jld2`): full 3-D temperature, salinity, velocity (and TKE if a `:e` tracer is present).
+
+`mode` selects the reduction:
+- `:averaged` uses `Oceananigans.AveragedTimeInterval(interval)` (time-averaged fields).
+- `:instantaneous` uses `Oceananigans.TimeInterval(interval)` (snapshots).
+
+!!! note
+    `tauuo`, `tauvo`, `hfds`, and `wfo` here are the *combined* surface fluxes on the ocean (atmosphere + sea-ice contributions),
+    not the atmosphere-only contribution. Sensible/latent heat (`hfss`/`hfls`) are not stored as standalone fields in CMIP —
+    they are folded into `oc_flux_T` inside `update_turbulent_fluxes!` — and are therefore omitted.
+"""
+function add_ocean_diagnostics!(
+    ocean_sim::OceananigansSimulation;
+    output_dir = joinpath("output_active", "clima_ocean"),
+    interval = Dates.Day(1),
+    mode = :averaged,
+    filename_prefix = "ocean",
+    file_splitting_interval = Dates.Day(15),
+)
+    ocean = ocean_sim.ocean
+    Nz = size(ocean.model.grid, 3)
+    file_splitting = OC.TimeInterval(file_splitting_interval)
+
+    T, S = ocean.model.tracers.T, ocean.model.tracers.S
+    u, v, w = ocean.model.velocities
+    η = ocean.model.free_surface.displacement
+
+    τx = surface_flux(u)
+    τy = surface_flux(v)
+    JT = surface_flux(T)
+    Js = surface_flux(S)
+
+    # Seed FieldStatus with the clock time so computed fields adopt the clock's time type
+    computed_field_status = OC.Fields.FieldStatus(ocean.model.clock.time)
+
+    surface_indices = (:, :, Nz)
+    sea_surface_temperature = OC.Field(T; indices = surface_indices)
+    sea_surface_salinity = OC.Field(S; indices = surface_indices)
+    sea_surface_zonal_velocity = OC.Field(u; indices = surface_indices)
+    sea_surface_meridional_velocity = OC.Field(v; indices = surface_indices)
+    sea_surface_temperature_squared =
+        OC.Field(T * T; indices = surface_indices, status = computed_field_status)
+    sea_surface_salinity_squared =
+        OC.Field(S * S; indices = surface_indices, status = computed_field_status)
+    sea_surface_height_squared = OC.Field(η * η; status = computed_field_status)
+
+    surface_outputs = Dict{Symbol, Any}(
+        :sea_surface_temperature => sea_surface_temperature,
+        :sea_surface_salinity => sea_surface_salinity,
+        :sea_surface_height => η,
+        :sea_surface_zonal_velocity => sea_surface_zonal_velocity,
+        :sea_surface_meridional_velocity => sea_surface_meridional_velocity,
+        :sea_surface_temperature_squared => sea_surface_temperature_squared,
+        :sea_surface_salinity_squared => sea_surface_salinity_squared,
+        :sea_surface_height_squared => sea_surface_height_squared,
+        :zonal_wind_stress => τx,
+        :meridional_wind_stress => τy,
+        :surface_heat_flux => JT,
+        :surface_salinity_flux => Js,
+    )
+
+    schedule = diagnostic_schedule(mode, interval)
+
+    ocean.output_writers[:surface_diagnostics] = OC.JLD2Writer(
+        ocean.model,
+        surface_outputs;
+        schedule,
+        filename = joinpath(output_dir, filename_prefix * "_surface"),
+        file_splitting,
+        overwrite_existing = true,
+    )
+
+    field_outputs = Dict{Symbol, Any}(
+        :temperature => T,
+        :salinity => S,
+        :zonal_velocity => u,
+        :meridional_velocity => v,
+        :vertical_velocity => w,
+    )
+
+    if haskey(ocean.model.tracers, :e)
+        field_outputs[:turbulent_kinetic_energy] = ocean.model.tracers.e
+    end
+
+    ocean.output_writers[:field_diagnostics] = OC.JLD2Writer(
+        ocean.model,
+        field_outputs;
+        schedule,
+        filename = joinpath(output_dir, filename_prefix * "_fields"),
+        file_splitting,
+        overwrite_existing = true,
+    )
+
+    @info "Ocean diagnostics attached:" *
+          " surface ($(length(surface_outputs)) fields, $mode, every $interval)," *
+          " 3-D ($(length(field_outputs)) fields, $mode, every $interval)"
+
+    return nothing
+end

ext/ClimaCouplerCMIPExt/oceananigans.jl

@@ -75,6 +75,8 @@ function OceananigansSimulation(
     comms_ctx = ClimaComms.context(),
     coupled_param_dict = CP.create_toml_dict(FT),
     progress_interval = nothing,
+    ocean_diagnostic_interval = "1days",
+    ocean_diagnostic_mode = :averaged,
     extra_kwargs...,
 ) where {FT}
     arch = comms_ctx.device isa ClimaComms.CUDADevice ? OC.GPU() : OC.CPU()
@@ -280,60 +282,30 @@ function OceananigansSimulation(
         coare3_roughness_params,
     )
 
-    # Before version 0.96.22, the NetCDFWriter was broken on GPU
-    if arch isa OC.CPU
-        # Save all tracers and velocities to a NetCDF file at daily frequency
-        outputs = merge(ocean.model.tracers, ocean.model.velocities)
-        surface_writer = OC.NetCDFWriter(
-            ocean.model,
-            outputs;
-            schedule = OC.TimeInterval(86400), # Daily output
-            filename = joinpath(output_dir, "ocean_diagnostics.nc"),
-            indices = (:, :, grid.Nz),
-            overwrite_existing = true,
-            array_type = Array{FT},
-        )
-        free_surface_writer = OC.NetCDFWriter(
-            ocean.model,
-            (; displacement = ocean.model.free_surface.displacement);
-            schedule = OC.TimeInterval(3600), # hourly snapshots
-            filename = joinpath(output_dir, "ocean_free_surface.nc"),
-            overwrite_existing = true,
-            array_type = Array{FT},
-        )
-        Tflux = ocean.model.tracers.T.boundary_conditions.top.condition
-        Sflux = ocean.model.tracers.S.boundary_conditions.top.condition
-        uflux = ocean.model.velocities.u.boundary_conditions.top.condition
-        vflux = ocean.model.velocities.v.boundary_conditions.top.condition
-        fluxes_writer = OC.NetCDFWriter(
-            ocean.model,
-            (; Tflux, Sflux, uflux, vflux);
-            schedule = OC.TimeInterval(3600), # hourly snapshots
-            filename = joinpath(output_dir, "ocean_fluxes.nc"),
-            overwrite_existing = true,
-            array_type = Array{FT},
-        )
-
-        ocean.output_writers[:surface] = surface_writer
-        ocean.output_writers[:free_surface] = free_surface_writer
-        ocean.output_writers[:fluxes] = fluxes_writer
-    end
-
     # Initialize with 0 ice concentration; this will be updated in `resolve_area_fractions!`
     # if the ocean is coupled to a non-prescribed sea ice model.
     ice_concentration = OC.Field{OC.Center, OC.Center, Nothing}(grid)
 
     # Create a dummy area fraction that will get overwritten in `update_surface_fractions!`
     area_fraction = ones(boundary_space)
 
-    return OceananigansSimulation(
+    sim = OceananigansSimulation(
         ocean,
         area_fraction,
         ocean_properties,
         remapping,
         ice_concentration,
         model_Δt,
     )
+
+    add_ocean_diagnostics!(
+        sim;
+        output_dir,
+        interval = TimeManager.time_to_period(ocean_diagnostic_interval),
+        mode = ocean_diagnostic_mode,
+    )
+
+    return sim
 end
 
 """