Move NetCDF output attribute methods from extension to Models module

Project.toml

@@ -1,6 +1,6 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.106.3"
+version = "0.106.4"
 authors = ["Climate Modeling Alliance and contributors"]
 
 [deps]

ext/OceananigansNCDatasetsExt/OceananigansNCDatasetsExt.jl

@@ -17,13 +17,11 @@ using NCDatasets: AbstractDataset
 using Dates: AbstractTime, UTC, now, DateTime
 using Printf: @sprintf
 using OrderedCollections: OrderedDict
-using SeawaterPolynomials: BoussinesqEquationOfState
 using Statistics: mean
 
 using Oceananigans: initialize!, prettytime, pretty_filesize, AbstractModel
 using Oceananigans.AbstractOperations: KernelFunctionOperation, AbstractOperation
 using Oceananigans.Architectures: CPU, GPU, on_architecture
-using Oceananigans.BuoyancyFormulations: BuoyancyForce, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.Fields
 using Oceananigans.Fields: set!, Reduction, reduced_dimensions, reduced_location, location, indices
 using Oceananigans.Grids:
@@ -35,7 +33,7 @@ using Oceananigans.Grids:
 using Oceananigans.ImmersedBoundaries:
     ImmersedBoundaryGrid, GridFittedBottom, GFBIBG, GridFittedBoundary, PartialCellBottom, PCBIBG,
     CenterImmersedCondition, InterfaceImmersedCondition
-using Oceananigans.Models: ShallowWaterModel, LagrangianParticles
+using Oceananigans.Models: LagrangianParticles
 using Oceananigans.OutputReaders:
     InMemoryFTS,
     time_indices,
@@ -79,8 +77,6 @@ import Oceananigans.OutputWriters:
 
 const c = Center()
 const f = Face()
-const BoussinesqSeawaterBuoyancy = SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S} where {FT, T, S}
-const BuoyancyBoussinesqEOSModel = BuoyancyForce{<:BoussinesqSeawaterBuoyancy, g} where {g}
 
 #####
 ##### Include scripts

ext/OceananigansNCDatasetsExt/netcdf_writer.jl

@@ -8,6 +8,8 @@
 # grid metrics.
 #
 
+using Oceananigans.OutputWriters: add_schedule_metadata!, default_output_attributes
+
 #####
 ##### Extend defVar to be able to write fields to NetCDF directly
 #####
@@ -59,86 +61,6 @@ function add_location_attribute!(attrib, fd::AbstractField)
     return merge(loc_attrib, attrib)
 end
 
-#####
-##### Variable attributes
-#####
-
-default_velocity_attributes(::RectilinearGrid) = Dict(
-    "u" => Dict("long_name" => "Velocity in the +x-direction.", "units" => "m/s"),
-    "v" => Dict("long_name" => "Velocity in the +y-direction.", "units" => "m/s"),
-    "w" => Dict("long_name" => "Velocity in the +z-direction.", "units" => "m/s"))
-
-default_velocity_attributes(::LatitudeLongitudeGrid) = Dict(
-    "u"            => Dict("long_name" => "Velocity in the zonal direction (+ = east).", "units" => "m/s"),
-    "v"            => Dict("long_name" => "Velocity in the meridional direction (+ = north).", "units" => "m/s"),
-    "w"            => Dict("long_name" => "Velocity in the vertical direction (+ = up).", "units" => "m/s"),
-    "displacement" => Dict("long_name" => "Sea surface height displacement", "units" => "m"))
-
-default_velocity_attributes(ibg::ImmersedBoundaryGrid) = default_velocity_attributes(ibg.underlying_grid)
-
-default_tracer_attributes(::Nothing) = Dict()
-
-default_tracer_attributes(::BuoyancyForce{<:BuoyancyTracer}) = Dict("b" => Dict("long_name" => "Buoyancy", "units" => "m/s²"))
-
-default_tracer_attributes(::BuoyancyForce{<:SeawaterBuoyancy{FT, <:LinearEquationOfState}}) where FT = Dict(
-    "T" => Dict("long_name" => "Temperature", "units" => "°C"),
-    "S" => Dict("long_name" => "Salinity",    "units" => "practical salinity unit (psu)"))
-
-default_tracer_attributes(::BuoyancyBoussinesqEOSModel) = Dict("T" => Dict("long_name" => "Conservative temperature", "units" => "°C"),
-                                                               "S" => Dict("long_name" => "Absolute salinity",        "units" => "g/kg"))
-
-function default_output_attributes(model)
-    velocity_attrs = default_velocity_attributes(model.grid)
-    buoyancy = model isa ShallowWaterModel ? nothing : model.buoyancy
-    tracer_attrs = default_tracer_attributes(buoyancy)
-    return merge(velocity_attrs, tracer_attrs)
-end
-
-#####
-##### Saving schedule metadata as global attributes
-#####
-
-add_schedule_metadata!(attributes, schedule) = nothing
-
-function add_schedule_metadata!(global_attributes, schedule::IterationInterval)
-    global_attributes["schedule"] = "IterationInterval"
-    global_attributes["interval"] = schedule.interval
-    global_attributes["output iteration interval"] = "Output was saved every $(schedule.interval) iteration(s)."
-
-    return nothing
-end
-
-function add_schedule_metadata!(global_attributes, schedule::TimeInterval)
-    global_attributes["schedule"] = "TimeInterval"
-    global_attributes["interval"] = schedule.interval
-    global_attributes["output time interval"] = "Output was saved every $(prettytime(schedule.interval))."
-
-    return nothing
-end
-
-function add_schedule_metadata!(global_attributes, schedule::WallTimeInterval)
-    global_attributes["schedule"] = "WallTimeInterval"
-    global_attributes["interval"] = schedule.interval
-    global_attributes["output time interval"] =
-        "Output was saved every $(prettytime(schedule.interval))."
-
-    return nothing
-end
-
-function add_schedule_metadata!(global_attributes, schedule::AveragedTimeInterval)
-    global_attributes["schedule"] = "AveragedTimeInterval"
-    global_attributes["interval"] = schedule.interval
-    global_attributes["output time interval"] = "Output was time-averaged and saved every $(prettytime(schedule.interval))."
-
-    global_attributes["time_averaging_window"] = schedule.window
-    global_attributes["time averaging window"] = "Output was time averaged with a window size of $(prettytime(schedule.window))"
-
-    global_attributes["time_averaging_stride"] = schedule.stride
-    global_attributes["time averaging stride"] = "Output was time averaged with a stride of $(schedule.stride) iteration(s) within the time averaging window."
-
-    return nothing
-end
-
 #####
 ##### NetCDFWriter constructor
 #####

src/Models/Models.jl

@@ -214,6 +214,9 @@ checkpointer_address(::HydrostaticFreeSurfaceModel) = "HydrostaticFreeSurfaceMod
 
 default_included_properties(::OceananigansModels) = [:grid]
 
+# Specialized output attributes for velocity and tracer fields
+include("output_attributes.jl")
+
 # Implementation of diagnostics applicable to both `NonhydrostaticModel` and `HydrostaticFreeSurfaceModel`
 include("seawater_density.jl")
 include("buoyancy_operation.jl")

src/Models/output_attributes.jl

@@ -0,0 +1,57 @@
+using Oceananigans.Grids: RectilinearGrid, LatitudeLongitudeGrid
+using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
+using Oceananigans.BuoyancyFormulations: BuoyancyForce, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
+
+using SeawaterPolynomials: BoussinesqEquationOfState
+
+import Oceananigans.OutputWriters: default_output_attributes
+
+const BoussinesqSeawaterBuoyancy = SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S} where {FT, T, S}
+const BuoyancyBoussinesqEOSModel = BuoyancyForce{<:BoussinesqSeawaterBuoyancy, g} where {g}
+
+#####
+##### Specialized velocity attributes
+#####
+
+default_velocity_attributes(::RectilinearGrid) = Dict(
+    "u" => Dict("long_name" => "Velocity in the +x-direction.", "units" => "m/s"),
+    "v" => Dict("long_name" => "Velocity in the +y-direction.", "units" => "m/s"),
+    "w" => Dict("long_name" => "Velocity in the +z-direction.", "units" => "m/s"))
+
+default_velocity_attributes(::LatitudeLongitudeGrid) = Dict(
+    "u"            => Dict("long_name" => "Velocity in the zonal direction (+ = east).", "units" => "m/s"),
+    "v"            => Dict("long_name" => "Velocity in the meridional direction (+ = north).", "units" => "m/s"),
+    "w"            => Dict("long_name" => "Velocity in the vertical direction (+ = up).", "units" => "m/s"),
+    "displacement" => Dict("long_name" => "Sea surface height displacement", "units" => "m"))
+
+default_velocity_attributes(ibg::ImmersedBoundaryGrid) = default_velocity_attributes(ibg.underlying_grid)
+
+#####
+##### Specialized tracer attributes
+#####
+
+default_tracer_attributes(::Nothing) = Dict{String, Any}()
+
+default_tracer_attributes(::BuoyancyForce{<:BuoyancyTracer}) = Dict("b" => Dict("long_name" => "Buoyancy", "units" => "m/s²"))
+
+default_tracer_attributes(::BuoyancyForce{<:SeawaterBuoyancy{FT, <:LinearEquationOfState}}) where FT = Dict(
+    "T" => Dict("long_name" => "Temperature", "units" => "°C"),
+    "S" => Dict("long_name" => "Salinity",    "units" => "practical salinity unit (psu)"))
+
+default_tracer_attributes(::BuoyancyBoussinesqEOSModel) = Dict(
+    "T" => Dict("long_name" => "Conservative temperature", "units" => "°C"),
+    "S" => Dict("long_name" => "Absolute salinity",        "units" => "g/kg"))
+
+#####
+##### Specialized default_output_attributes
+#####
+
+function default_output_attributes(model::Union{NonhydrostaticModel, HydrostaticFreeSurfaceModel})
+    velocity_attrs = default_velocity_attributes(model.grid)
+    tracer_attrs = default_tracer_attributes(model.buoyancy)
+    return merge(velocity_attrs, tracer_attrs)
+end
+
+function default_output_attributes(model::ShallowWaterModel)
+    return default_velocity_attributes(model.grid)
+end

src/OutputWriters/OutputWriters.jl

@@ -34,6 +34,7 @@ include("averaged_specified_times.jl")
 include("windowed_time_average.jl")
 include("output_construction.jl")
 include("jld2_writer.jl")
+include("output_attributes.jl")
 include("netcdf_writer.jl")
 include("checkpointer.jl")
 

src/OutputWriters/output_attributes.jl

@@ -0,0 +1,52 @@
+using Oceananigans.Utils: prettytime, TimeInterval, IterationInterval, WallTimeInterval
+
+#####
+##### Generic fallbacks for variable attributes
+#####
+
+default_output_attributes(model) = Dict{String, Any}()
+
+#####
+##### Saving schedule metadata as global attributes
+#####
+
+add_schedule_metadata!(attributes, schedule) = nothing
+
+function add_schedule_metadata!(global_attributes, schedule::IterationInterval)
+    global_attributes["schedule"] = "IterationInterval"
+    global_attributes["interval"] = schedule.interval
+    global_attributes["output iteration interval"] = "Output was saved every $(schedule.interval) iteration(s)."
+
+    return nothing
+end
+
+function add_schedule_metadata!(global_attributes, schedule::TimeInterval)
+    global_attributes["schedule"] = "TimeInterval"
+    global_attributes["interval"] = schedule.interval
+    global_attributes["output time interval"] = "Output was saved every $(prettytime(schedule.interval))."
+
+    return nothing
+end
+
+function add_schedule_metadata!(global_attributes, schedule::WallTimeInterval)
+    global_attributes["schedule"] = "WallTimeInterval"
+    global_attributes["interval"] = schedule.interval
+    global_attributes["output time interval"] =
+        "Output was saved every $(prettytime(schedule.interval))."
+
+    return nothing
+end
+
+function add_schedule_metadata!(global_attributes, schedule::AveragedTimeInterval)
+    global_attributes["schedule"] = "AveragedTimeInterval"
+    global_attributes["interval"] = schedule.interval
+    global_attributes["output time interval"] = "Output was time-averaged and saved every $(prettytime(schedule.interval))."
+
+    global_attributes["time_averaging_window"] = schedule.window
+    global_attributes["time averaging window"] = "Output was time averaged with a window size of $(prettytime(schedule.window))"
+
+    global_attributes["time_averaging_stride"] = schedule.stride
+    global_attributes["time averaging stride"] = "Output was time averaged with a stride of $(schedule.stride) iteration(s) within the time averaging window."
+
+    return nothing
+end