diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl index 156c3b17cb..c4f27997bf 100644 --- a/ext/OceananigansNCDatasetsExt.jl +++ b/ext/OceananigansNCDatasetsExt.jl @@ -46,7 +46,8 @@ import Oceananigans.OutputWriters: convert_for_netcdf, materialize_from_netcdf, reconstruct_grid, - trilocation_dim_name + trilocation_dim_name, + dimension_name_generator_free_surface const c = Center() const f = Face() @@ -57,21 +58,84 @@ const BuoyancyBoussinesqEOSModel = BuoyancyForce{<:BoussinesqSeawaterBuoyancy, g ##### Extend defVar to be able to write fields to NetCDF directly ##### +""" + squeeze_data(fd::AbstractField; array_type=Array{eltype(fd)}) + +Returns the data of the field with the any dimensions where location is Nothing squeezed. For example: +```Julia +infil> grid = RectilinearGrid(size=(2,3,4), extent=(1,1,1)) +2×3×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo +├── Periodic x ∈ [0.0, 1.0) regularly spaced with Δx=0.5 +├── Periodic y ∈ [0.0, 1.0) regularly spaced with Δy=0.333333 +└── Bounded z ∈ [-1.0, 0.0] regularly spaced with Δz=0.25 + +infil> c = Field{Center, Center, Nothing}(grid) +2×3×1 Field{Center, Center, Nothing} reduced over dims = (3,) on RectilinearGrid on CPU +├── grid: 2×3×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 2×3×3 halo +├── boundary conditions: FieldBoundaryConditions +│ └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: Nothing, top: Nothing, immersed: Nothing +└── data: 6×9×1 OffsetArray(::Array{Float64, 3}, -1:4, -2:6, 1:1) with eltype Float64 with indices -1:4×-2:6×1:1 + └── max=0.0, min=0.0, mean=0.0 + +infil> interior(c) |> size +(2, 3, 1) + +infil> squeeze_data(c) +2×3 Matrix{Float64}: + 0.0 0.0 0.0 + 0.0 0.0 0.0 + +infil> squeeze_data(c) |> size +(2, 3) +``` + +Note that this will only remove (squeeze) singleton dimensions. +""" +function squeeze_data(fd::AbstractField, field_data; array_type=Array{eltype(fd)}) + reduced_dims = effective_reduced_dimensions(fd) + field_data_cpu = array_type(field_data) # Need to convert to the array type of the field + + indices = Any[:, :, :] + for i in 1:3 + if i ∈ reduced_dims + indices[i] = 1 + end + end + return getindex(field_data_cpu, indices...) +end + +squeeze_data(func, func_data; kwargs...) = func_data +squeeze_data(wta::WindowedTimeAverage{<:AbstractField}, data; kwargs...) = squeeze_data(wta.operand, data; kwargs...) +squeeze_data(fd::AbstractField; kwargs...) = squeeze_data(fd, parent(fd); kwargs...) + +squeeze_data(fd::WindowedTimeAverage{<:AbstractField}; kwargs...) = squeeze_data(fd.operand; kwargs...) + defVar(ds, name, op::AbstractOperation; kwargs...) = defVar(ds, name, Field(op); kwargs...) defVar(ds, name, op::Reduction; kwargs...) = defVar(ds, name, Field(op); kwargs...) -function defVar(ds, name, field::AbstractField; +function defVar(ds, field_name, fd::AbstractField; + array_type=Array{eltype(fd)}, time_dependent=false, + with_halos=false, dimension_name_generator = trilocation_dim_name, + dimension_type=Float64, + write_data=true, kwargs...) - field_cpu = on_architecture(CPU(), field) # Need to bring field to CPU in order to write it to NetCDF - field_data = Array{eltype(field)}(field_cpu) - dims = field_dimensions(field, dimension_name_generator) - all_dims = time_dependent ? (dims..., "time") : dims - # Validate that all dimensions exist and match the field - create_field_dimensions!(ds, field, all_dims, dimension_name_generator) - defVar(ds, name, field_data, all_dims; kwargs...) + # effective_dim_names are the dimensions that will be used to write the field data (excludes reduced and dimensions where location is Nothing) + effective_dim_names = create_field_dimensions!(ds, fd, dimension_name_generator; time_dependent, with_halos, array_type, dimension_type) + + # Write the data to the NetCDF file (or don't, but still create the space for it there) + if write_data + # Squeeze the data to remove dimensions where location is Nothing and add a time dimension if the field is time-dependent + constructed_fd = construct_output(fd, fd.grid, (:, :, :), with_halos) + squeezed_field_data = squeeze_data(constructed_fd; array_type) + squeezed_reshaped_field_data = time_dependent ? reshape(squeezed_field_data, size(squeezed_field_data)..., 1) : squeezed_field_data + + defVar(ds, field_name, squeezed_reshaped_field_data, effective_dim_names; kwargs...) + else + defVar(ds, field_name, eltype(array_type), effective_dim_names; kwargs...) + end end ##### @@ -79,30 +143,35 @@ end ##### """ - create_field_dimensions!(ds, field::AbstractField, all_dims, dimension_name_generator) + create_field_dimensions!(ds, fd::AbstractField, dimension_name_generator; time_dependent=false, with_halos=false, array_type=Array{eltype(fd)}) -Creates all dimensions for the given `field` in the NetCDF dataset `ds`. If the dimensions -already exist, they are validated to match the expected dimensions for the given `field`. +Creates all dimensions for the given field `fd` in the NetCDF dataset `ds`. If the dimensions +already exist, they are validated to match the expected dimensions. Arguments: - `ds`: NetCDF dataset -- `field`: AbstractField being written -- `all_dims`: Tuple of dimension names to create/validate +- `fd`: AbstractField being written +- `dim_names`: Tuple of dimension names to create/validate - `dimension_name_generator`: Function to generate dimension names """ -function create_field_dimensions!(ds, field::AbstractField, all_dims, dimension_name_generator) - dimension_attributes = default_dimension_attributes(field.grid, dimension_name_generator) - spatial_dims = all_dims[1:end-(("time" in all_dims) ? 1 : 0)] +function create_field_dimensions!(ds, fd::AbstractField, dimension_name_generator; time_dependent=false, with_halos=false, array_type=Array{eltype(fd)}, dimension_type=Float64) + # Assess and create the dimensions for the field - spatial_dims_dict = Dict(dim_name => dim_data for (dim_name, dim_data) in zip(spatial_dims, nodes(field))) - create_spatial_dimensions!(ds, spatial_dims_dict, dimension_attributes; array_type=Array{eltype(field)}) + dimension_attributes = default_dimension_attributes(fd.grid, dimension_name_generator) + spatial_dim_names = field_dimensions(fd, dimension_name_generator) + spatial_dim_data = nodes(fd; with_halos) + + # Create dictionary of spatial dimensions and their data. Using OrderedDict to ensure the order of the dimensions is preserved. + spatial_dim_names_dict = OrderedDict(spatial_dim_name => spatial_dim_array for (spatial_dim_name, spatial_dim_array) in zip(spatial_dim_names, spatial_dim_data)) + effective_spatial_dim_names = create_spatial_dimensions!(ds, spatial_dim_names_dict, dimension_attributes; dimension_type) # Create time dimension if needed - if "time" in all_dims && "time" ∉ keys(ds.dim) - create_time_dimension!(ds) + if time_dependent + "time" ∉ keys(ds.dim) && create_time_dimension!(ds, dimension_type=dimension_type) + return (effective_spatial_dim_names..., "time") # Add "time" dimension if the field is time-dependent + else + return effective_spatial_dim_names end - - return nothing end ##### @@ -124,13 +193,11 @@ function collect_dim(ξ, ℓ, T, N, H, inds, with_halos) end end -function create_time_dimension!(dataset; attrib=nothing) +function create_time_dimension!(dataset; attrib=nothing, dimension_type=Float64) if "time" ∉ keys(dataset.dim) # Create an unlimited dimension "time" - # Time should always be Float64 to be extra safe from rounding errors. - # See: https://github.com/CliMA/Oceananigans.jl/issues/3056 defDim(dataset, "time", Inf) - defVar(dataset, "time", Float64, ("time",), attrib=attrib) + defVar(dataset, "time", dimension_type, ("time",), attrib=attrib) end end @@ -144,47 +211,47 @@ their coordinate values. Each dimension variable has itself as its sole dimensio against provided arrays if they do exist. An error is thrown if the dimension already exists but is different from the provided array. """ -function create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=Array{Float32}, kwargs...) +function create_spatial_dimensions!(dataset, dims, attributes_dict; dimension_type=Float64, kwargs...) + effective_dim_names = [] for (i, (dim_name, dim_array)) in enumerate(dims) + dim_array isa Nothing && continue # Don't create anything if dim_array is Nothing + dim_name == "" && continue # Don't create anything if dim_name is an empty string + push!(effective_dim_names, dim_name) + + dim_array = dimension_type.(dim_array) # Transform dim_array to the correct float type if dim_name ∉ keys(dataset.dim) # Create missing dimension - defVar(dataset, dim_name, array_type(dim_array), (dim_name,), attrib=attributes_dict[dim_name]; kwargs...) + defVar(dataset, dim_name, dim_array, (dim_name,), attrib=attributes_dict[dim_name]; kwargs...) else # Validate existing dimension - if dataset[dim_name] != dim_array + dataset_dim_array = collect(dataset[dim_name]) + if dataset_dim_array != collect(dim_array) throw(ArgumentError("Dimension '$dim_name' already exists in dataset but is different from expected.\n" * - " Actual: $(dataset[dim_name]) (length=$(length(dataset[dim_name])))\n" * + " Actual: $(dataset_dim_array) (length=$(length(dataset_dim_array)))\n" * " Expected: $(dim_array) (length=$(length(dim_array)))")) end end end + return tuple(effective_dim_names...) end """ effective_reduced_dimensions(field) Return dimensions that are effectively reduced, considering both location-based reduction -(e.g. a `Nothing` location) and index-based reduction at boundaries (e.g. free surface -height fields with :, :, Nz+1:Nz+1 indices). -``` +(e.g. a `Nothing` location) and grid topology (i.e. a `Flat` topology is considered a reduction). """ function effective_reduced_dimensions(field) loc_reduced = reduced_dimensions(field) - idx_reduced = () - inds = indices(field) - grid_size = size(field.grid) - - for (dim, ind) in enumerate(inds) - if ind isa UnitRange && length(ind) == 1 - index_value = first(ind) - if index_value > grid_size[dim] - idx_reduced = (idx_reduced..., dim) - end + topo_reduced = [] + for (dim, topo) in enumerate(topology(field)) + if topo == Flat + push!(topo_reduced, dim) end end - all_reduced = (loc_reduced..., idx_reduced...) + all_reduced = (loc_reduced..., topo_reduced...) return Tuple(unique(all_reduced)) end @@ -453,45 +520,33 @@ end ##### Mapping outputs/fields to dimensions ##### -function field_dimensions(field::AbstractField, grid::RectilinearGrid, dim_name_generator) - LX, LY, LZ = location(field) +function field_dimensions(fd::AbstractField, grid::RectilinearGrid, dim_name_generator) + LX, LY, LZ = location(fd) TX, TY, TZ = topology(grid) - eff_reduced_dims = effective_reduced_dimensions(field) - - x_dim_name = (1 ∈ eff_reduced_dims || TX == Flat) ? "" : dim_name_generator("x", grid, LX(), nothing, nothing, Val(:x)) - y_dim_name = (2 ∈ eff_reduced_dims || TY == Flat) ? "" : dim_name_generator("y", grid, nothing, LY(), nothing, Val(:y)) - z_dim_name = (3 ∈ eff_reduced_dims || TZ == Flat) ? "" : dim_name_generator("z", grid, nothing, nothing, LZ(), Val(:z)) - - x_dim_name = isempty(x_dim_name) ? tuple() : tuple(x_dim_name) - y_dim_name = isempty(y_dim_name) ? tuple() : tuple(y_dim_name) - z_dim_name = isempty(z_dim_name) ? tuple() : tuple(z_dim_name) + x_dim_name = LX == Nothing ? "" : dim_name_generator("x", grid, LX(), nothing, nothing, Val(:x)) + y_dim_name = LY == Nothing ? "" : dim_name_generator("y", grid, nothing, LY(), nothing, Val(:y)) + z_dim_name = LZ == Nothing ? "" : dim_name_generator("z", grid, nothing, nothing, LZ(), Val(:z)) - return tuple(x_dim_name..., y_dim_name..., z_dim_name...) + return tuple(x_dim_name, y_dim_name, z_dim_name) end -function field_dimensions(field::AbstractField, grid::LatitudeLongitudeGrid, dim_name_generator) - LΛ, LΦ, LZ = location(field) +function field_dimensions(fd::AbstractField, grid::LatitudeLongitudeGrid, dim_name_generator) + LΛ, LΦ, LZ = location(fd) TΛ, TΦ, TZ = topology(grid) - eff_reduced_dims = effective_reduced_dimensions(field) + λ_dim_name = LΛ == Nothing ? "" : dim_name_generator("λ", grid, LΛ(), nothing, nothing, Val(:x)) + φ_dim_name = LΦ == Nothing ? "" : dim_name_generator("φ", grid, nothing, LΦ(), nothing, Val(:y)) + z_dim_name = LZ == Nothing ? "" : dim_name_generator("z", grid, nothing, nothing, LZ(), Val(:z)) - λ_dim_name = (1 ∈ eff_reduced_dims || TΛ == Flat) ? "" : dim_name_generator("λ", grid, LΛ(), nothing, nothing, Val(:x)) - φ_dim_name = (2 ∈ eff_reduced_dims || TΦ == Flat) ? "" : dim_name_generator("φ", grid, nothing, LΦ(), nothing, Val(:y)) - z_dim_name = (3 ∈ eff_reduced_dims || TZ == Flat) ? "" : dim_name_generator("z", grid, nothing, nothing, LZ(), Val(:z)) - - λ_dim_name = isempty(λ_dim_name) ? tuple() : tuple(λ_dim_name) - φ_dim_name = isempty(φ_dim_name) ? tuple() : tuple(φ_dim_name) - z_dim_name = isempty(z_dim_name) ? tuple() : tuple(z_dim_name) - - return tuple(λ_dim_name..., φ_dim_name..., z_dim_name...) + return tuple(λ_dim_name, φ_dim_name, z_dim_name) end -field_dimensions(field::AbstractField, grid::ImmersedBoundaryGrid, dim_name_generator) = - field_dimensions(field, grid.underlying_grid, dim_name_generator) +field_dimensions(fd::AbstractField, grid::ImmersedBoundaryGrid, dim_name_generator) = + field_dimensions(fd, grid.underlying_grid, dim_name_generator) -field_dimensions(field::AbstractField, dim_name_generator) = - field_dimensions(field, field.grid, dim_name_generator) +field_dimensions(fd::AbstractField, dim_name_generator) = + field_dimensions(fd, fd.grid, dim_name_generator) ##### ##### Dimension attributes @@ -1070,7 +1125,8 @@ function NetCDFWriter(model::AbstractModel, outputs; deflatelevel = 0, part = 1, file_splitting = NoFileSplitting(), - dimension_name_generator = trilocation_dim_name) + dimension_name_generator = trilocation_dim_name, + dimension_type = Float64) if with_halos && indices != (:, :, :) throw(ArgumentError("If with_halos=true then you cannot pass indices: $indices")) @@ -1111,6 +1167,7 @@ function NetCDFWriter(model::AbstractModel, outputs; global_attributes = Dict{Any, Any}(global_attributes) dataset, outputs, schedule = initialize_nc_file(model, + grid, filepath, outputs, schedule, @@ -1123,7 +1180,8 @@ function NetCDFWriter(model::AbstractModel, outputs; include_grid_metrics, overwrite_existing, deflatelevel, - dimension_name_generator) + dimension_name_generator, + dimension_type) return NetCDFWriter(grid, filepath, @@ -1142,7 +1200,8 @@ function NetCDFWriter(model::AbstractModel, outputs; deflatelevel, part, file_splitting, - dimension_name_generator) + dimension_name_generator, + dimension_type) end ##### @@ -1150,6 +1209,7 @@ end ##### function initialize_nc_file(model, + grid, filepath, outputs, schedule, @@ -1162,9 +1222,9 @@ function initialize_nc_file(model, include_grid_metrics, overwrite_existing, deflatelevel, - dimension_name_generator) + dimension_name_generator, + dimension_type) - grid = model.grid mode = overwrite_existing ? "c" : "a" # Add useful metadata @@ -1206,8 +1266,8 @@ function initialize_nc_file(model, Dict("long_name" => "Time", "units" => "seconds since 2000-01-01 00:00:00") : Dict("long_name" => "Time", "units" => "seconds") - create_time_dimension!(dataset, attrib=time_attrib) - create_spatial_dimensions!(dataset, dims, output_attributes; deflatelevel=1, array_type) + create_time_dimension!(dataset, attrib=time_attrib, dimension_type=dimension_type) + create_spatial_dimensions!(dataset, dims, output_attributes; deflatelevel=1, dimension_type=dimension_type) time_independent_vars = Dict() @@ -1222,40 +1282,46 @@ function initialize_nc_file(model, end if !isempty(time_independent_vars) - for (name, output) in sort(collect(pairs(time_independent_vars)), by=first) + for (output_name, output) in sort(collect(pairs(time_independent_vars)), by=first) output = construct_output(output, grid, indices, with_halos) - attributes = haskey(output_attributes, name) ? output_attributes[name] : Dict() + attrib = haskey(output_attributes, output_name) ? output_attributes[output_name] : Dict() materialized = materialize_output(output, model) - define_output_variable!(dataset, + define_output_variable!(model, + dataset, materialized, - name, + output_name; array_type, deflatelevel, - attributes, + attrib, dimensions, filepath, # for better error messages dimension_name_generator, - false) # time_dependent = false + time_dependent = false, + with_halos, + dimension_type) - save_output!(dataset, output, model, name, array_type) + save_output!(dataset, output, model, output_name, array_type) end end - for (name, output) in sort(collect(pairs(outputs)), by=first) - attributes = haskey(output_attributes, name) ? output_attributes[name] : Dict() + for (output_name, output) in sort(collect(pairs(outputs)), by=first) + attrib = haskey(output_attributes, output_name) ? output_attributes[output_name] : Dict() materialized = materialize_output(output, model) - define_output_variable!(dataset, + define_output_variable!(model, + dataset, materialized, - name, + output_name; array_type, deflatelevel, - attributes, + attrib, dimensions, filepath, # for better error messages dimension_name_generator, - true) # time_dependent = true) + time_dependent = true, + with_halos, + dimension_type) end sync(dataset) @@ -1267,6 +1333,7 @@ function initialize_nc_file(model, end initialize_nc_file(ow::NetCDFWriter, model) = initialize_nc_file(model, + ow.grid, ow.filepath, ow.outputs, ow.schedule, @@ -1279,7 +1346,8 @@ initialize_nc_file(ow::NetCDFWriter, model) = initialize_nc_file(model, ow.include_grid_metrics, ow.overwrite_existing, ow.deflatelevel, - ow.dimension_name_generator) + ow.dimension_name_generator, + ow.dimension_type) ##### ##### Variable definition @@ -1291,46 +1359,48 @@ materialize_output(particles::LagrangianParticles, model) = particles materialize_output(output::WindowedTimeAverage{<:AbstractField}, model) = output """ Defines empty variables for 'custom' user-supplied `output`. """ -function define_output_variable!(dataset, output, name, array_type, - deflatelevel, attrib, dimensions, filepath, - dimension_name_generator, time_dependent) +function define_output_variable!(model, dataset, output, output_name; array_type, + deflatelevel, attrib, dimension_name_generator, + time_dependent, with_halos, + dimensions, filepath, dimension_type=Float64) - if name ∉ keys(dimensions) - msg = string("dimensions[$name] for output $name=$(typeof(output)) into $filepath" * + if output_name ∉ keys(dimensions) + msg = string("dimensions[$output_name] for output $output_name=$(typeof(output)) into $filepath" * " must be provided when constructing NetCDFWriter") throw(ArgumentError(msg)) end - dims = dimensions[name] + dims = dimensions[output_name] FT = eltype(array_type) - defVar(dataset, name, FT, (dims..., "time"); deflatelevel, attrib) + all_dims = time_dependent ? (dims..., "time") : dims + defVar(dataset, output_name, FT, all_dims; deflatelevel, attrib) return nothing end """ Defines empty field variable. """ -function define_output_variable!(dataset, output::AbstractField, name, array_type, - deflatelevel, attrib, dimensions, filepath, - dimension_name_generator, time_dependent) - - dims = field_dimensions(output, dimension_name_generator) - FT = eltype(array_type) - - all_dims = time_dependent ? (dims..., "time") : dims - - defVar(dataset, name, FT, all_dims; deflatelevel, attrib) - +function define_output_variable!(model, dataset, output::AbstractField, output_name; array_type, + deflatelevel, attrib, dimension_name_generator, + time_dependent, with_halos, + dimensions, filepath, dimension_type=Float64) + + # If the output is the free surface, we need to handle it differently since it will be writen as a 3D array with a singleton dimension for the z-coordinate + if output_name == "η" && output == view(model.free_surface.η, output.indices...) + local default_dimension_name_generator = dimension_name_generator + dimension_name_generator = (var_name, grid, LX, LY, LZ, dim) -> dimension_name_generator_free_surface(default_dimension_name_generator, var_name, grid, LX, LY, LZ, dim) + end + defVar(dataset, output_name, output; array_type, time_dependent, with_halos, dimension_name_generator, deflatelevel, attrib, dimension_type, write_data=false) return nothing end """ Defines empty field variable for `WindowedTimeAverage`s over fields. """ -define_output_variable!(dataset, output::WindowedTimeAverage{<:AbstractField}, args...) = - define_output_variable!(dataset, output.operand, args...) +define_output_variable!(model, dataset, output::WindowedTimeAverage{<:AbstractField}, output_name; kwargs...) = + define_output_variable!(model, dataset, output.operand, output_name; kwargs...) """ Defines empty variable for particle trackting. """ -function define_output_variable!(dataset, output::LagrangianParticles, name, array_type, - deflatelevel, args...) +function define_output_variable!(model, dataset, output::LagrangianParticles, output_name; array_type, + deflatelevel, kwargs...) particle_fields = eltype(output.properties) |> fieldnames .|> string T = eltype(array_type) @@ -1350,21 +1420,21 @@ Base.open(nc::NetCDFWriter) = NCDataset(nc.filepath, "a") Base.close(nc::NetCDFWriter) = close(nc.dataset) # Saving outputs with no time dependence (e.g. grid metrics) -function save_output!(ds, output, model, name, array_type) +function save_output!(ds, output, model, output_name, array_type) fetched = fetch_output(output, model) data = convert_output(fetched, array_type) - data = drop_output_dims(output, data) + data = squeeze_data(output, data) colons = Tuple(Colon() for _ in 1:ndims(data)) - ds[name][colons...] = data + ds[output_name][colons...] = data return nothing end # Saving time-dependent outputs -function save_output!(ds, output, model, ow, time_index, name) +function save_output!(ds, output, model, ow, time_index, output_name) data = fetch_and_convert_output(output, model, ow) - data = drop_output_dims(output, data) + data = squeeze_data(output, data) colons = Tuple(Colon() for _ in 1:ndims(data)) - ds[name][colons..., time_index:time_index] = data + ds[output_name][colons..., time_index:time_index] = data return nothing end @@ -1407,16 +1477,16 @@ function write_output!(ow::NetCDFWriter, model::AbstractModel) t0, sz0 = time_ns(), filesize(filepath) end - for (name, output) in ow.outputs + for (output_name, output) in ow.outputs # Time before computing this output. verbose && (t0′ = time_ns()) - save_output!(ds, output, model, ow, time_index, name) + save_output!(ds, output, model, ow, time_index, output_name) if verbose # Time after computing this output. t1′ = time_ns() - @info "Computing $name done: time=$(prettytime((t1′-t0′) / 1e9))" + @info "Computing $output_name done: time=$(prettytime((t1′-t0′) / 1e9))" end end @@ -1435,21 +1505,6 @@ function write_output!(ow::NetCDFWriter, model::AbstractModel) return nothing end -drop_output_dims(output, data) = data # fallback -drop_output_dims(output::WindowedTimeAverage{<:Field}, data) = drop_output_dims(output.operand, data) - -function drop_output_dims(field::Field, data) - eff_reduced_dims = effective_reduced_dimensions(field) - flat_dims = Tuple(i for (i, T) in enumerate(topology(field.grid)) if T == Flat) - dims = (eff_reduced_dims..., flat_dims...) - dims = Tuple(Set(dims)) # ensure dims are unique - - # Only drop dimensions that exist in the data and are size 1 - dims = filter(d -> d <= ndims(data) && size(data, d) == 1, dims) - - return isempty(dims) ? data : dropdims(data; dims=tuple(dims...)) -end - ##### ##### Show ##### diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl index 46c252306a..9597b4945c 100644 --- a/src/Grids/latitude_longitude_grid.jl +++ b/src/Grids/latitude_longitude_grid.jl @@ -632,6 +632,7 @@ end const F = Face const C = Center +const N = Nothing @inline function xnodes(grid::LLG, ℓx, ℓy; with_halos=false) λ = λnodes(grid, ℓx; with_halos=with_halos)' @@ -654,8 +655,10 @@ end @inline λnodes(grid::LLG, ℓx::F; with_halos=false, indices=Colon()) = view(_property(grid.λᶠᵃᵃ, ℓx, topology(grid, 1), grid.Nx, grid.Hx, with_halos), indices) @inline λnodes(grid::LLG, ℓx::C; with_halos=false, indices=Colon()) = view(_property(grid.λᶜᵃᵃ, ℓx, topology(grid, 1), grid.Nx, grid.Hx, with_halos), indices) +@inline λnodes(grid::LLG, ℓx::N; with_halos=false, indices=Colon()) = nothing @inline φnodes(grid::LLG, ℓy::F; with_halos=false, indices=Colon()) = view(_property(grid.φᵃᶠᵃ, ℓy, topology(grid, 2), grid.Ny, grid.Hy, with_halos), indices) @inline φnodes(grid::LLG, ℓy::C; with_halos=false, indices=Colon()) = view(_property(grid.φᵃᶜᵃ, ℓy, topology(grid, 2), grid.Ny, grid.Hy, with_halos), indices) +@inline φnodes(grid::LLG, ℓy::N; with_halos=false, indices=Colon()) = nothing # Flat topologies XFlatLLG = LatitudeLongitudeGrid{<:Any, Flat} diff --git a/src/OutputWriters/netcdf_writer.jl b/src/OutputWriters/netcdf_writer.jl index e1a8a90bcb..50333dc47b 100644 --- a/src/OutputWriters/netcdf_writer.jl +++ b/src/OutputWriters/netcdf_writer.jl @@ -62,9 +62,10 @@ trilocation_dim_name(var_name, grid, LX, LY, LZ, dim) = trilocation_dim_name(var_name, grid::ImmersedBoundaryGrid, args...) = trilocation_dim_name(var_name, grid.underlying_grid, args...) +dimension_name_generator_free_surface(dimension_name_generator, var_name, grid, LX, LY, LZ, dim) = dimension_name_generator(var_name, grid, LX, LY, LZ, dim) +dimension_name_generator_free_surface(dimension_name_generator, var_name, grid, LX, LY, LZ, dim::Val{:z}) = dimension_name_generator(var_name, grid, LX, LY, LZ, dim) * "_η" - -mutable struct NetCDFWriter{G, D, O, T, A, FS, DN} <: AbstractOutputWriter +mutable struct NetCDFWriter{G, D, O, T, A, FS, DN, DT} <: AbstractOutputWriter grid :: G filepath :: String dataset :: D @@ -83,6 +84,7 @@ mutable struct NetCDFWriter{G, D, O, T, A, FS, DN} <: AbstractOutputWriter part :: Int file_splitting :: FS dimension_name_generator :: DN + dimension_type :: DT end function NetCDFWriter(model, outputs; kw...) diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl index 94f7cc8600..28b75f4af7 100644 --- a/test/test_netcdf_writer.jl +++ b/test/test_netcdf_writer.jl @@ -13,6 +13,7 @@ using Oceananigans: Clock using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant using Oceananigans.OutputWriters: trilocation_dim_name using Oceananigans.Grids: ξname, ηname, rname, ξnodes, ηnodes +using Oceananigans.Fields: interpolate! function test_datetime_netcdf_output(arch) grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1)) @@ -119,14 +120,13 @@ function test_netcdf_grid_metrics_rectilinear(arch, FT) filepath_metrics_halos = "test_grid_metrics_rectilinear_halos_$(Arch)_$FT.nc" isfile(filepath_metrics_halos) && rm(filepath_metrics_halos) - simulation.output_writers[:with_metrics_and_halos] = - NetCDFWriter(model, fields(model), - filename = filepath_metrics_halos, - schedule = IterationInterval(1), - array_type = Array{FT}, - with_halos = true, - include_grid_metrics = true, - verbose = true) + simulation.output_writers[:with_metrics_and_halos] = NetCDFWriter(model, fields(model), + filename = filepath_metrics_halos, + schedule = IterationInterval(1), + array_type = Array{FT}, + with_halos = true, + include_grid_metrics = true, + verbose = true) filepath_metrics_nohalos = "test_grid_metrics_rectilinear_nohalos_$(Arch)_$FT.nc" isfile(filepath_metrics_nohalos) && rm(filepath_metrics_nohalos) @@ -176,19 +176,27 @@ function test_netcdf_grid_metrics_rectilinear(arch, FT) ds_mh = NCDataset(filepath_metrics_halos) @test haskey(ds_mh, "time") - @test eltype(ds_mh["time"]) == Float64 + @test eltype(ds_mh["time"]) == Float64 # All dimensions should be Float64 by default dims = ("x_faa", "x_caa", "y_afa", "y_aca", "z_aaf", "z_aac") metrics = ("Δx_faa", "Δx_caa", "Δy_afa", "Δy_aca", "Δy_afa", "Δy_aca") vars = ("u", "v", "w", "T", "S") - for var in (dims..., metrics..., vars...) + for var in (metrics..., vars...) @test haskey(ds_mh, var) @test haskey(ds_mh[var].attrib, "long_name") @test haskey(ds_mh[var].attrib, "units") @test eltype(ds_mh[var]) == FT end + for var in dims + @test haskey(ds_mh, var) + @test haskey(ds_mh[var].attrib, "long_name") + @test haskey(ds_mh[var].attrib, "units") + @test eltype(ds_mh[var]) == Float64 # All dimensions should be Float64 by default + end + + @test dimsize(ds_mh["time"]) == (time=Nt + 1,) @test dimsize(ds_mh[:x_faa]) == (x_faa=Nx + 2Hx,) @@ -218,15 +226,22 @@ function test_netcdf_grid_metrics_rectilinear(arch, FT) ds_m = NCDataset(filepath_metrics_nohalos) @test haskey(ds_m, "time") - @test eltype(ds_m["time"]) == Float64 + @test eltype(ds_m["time"]) == Float64 # All dimensions should be Float64 by default - for var in (dims..., metrics..., vars...) + for var in (metrics..., vars...) @test haskey(ds_m, var) @test haskey(ds_m[var].attrib, "long_name") @test haskey(ds_m[var].attrib, "units") @test eltype(ds_m[var]) == FT end + for var in dims + @test haskey(ds_m, var) + @test haskey(ds_m[var].attrib, "long_name") + @test haskey(ds_m[var].attrib, "units") + @test eltype(ds_m[var]) == Float64 # All dimensions should be Float64 by default + end + @test dimsize(ds_m[:x_faa]) == (x_faa=Nx,) @test dimsize(ds_m[:x_caa]) == (x_caa=Nx,) @test dimsize(ds_m[:y_afa]) == (y_afa=Ny + 1,) @@ -254,15 +269,22 @@ function test_netcdf_grid_metrics_rectilinear(arch, FT) ds_h = NCDataset(filepath_nometrics) @test haskey(ds_h, "time") - @test eltype(ds_h["time"]) == Float64 + @test eltype(ds_h["time"]) == Float64 # All dimensions should be Float64 by default - for var in (dims..., vars...) + for var in vars @test haskey(ds_h, var) @test haskey(ds_h[var].attrib, "long_name") @test haskey(ds_h[var].attrib, "units") @test eltype(ds_h[var]) == FT end + for var in dims + @test haskey(ds_h, var) + @test haskey(ds_h[var].attrib, "long_name") + @test haskey(ds_h[var].attrib, "units") + @test eltype(ds_h[var]) == Float64 # All dimensions should be Float64 by default + end + for metric in metrics @test !haskey(ds_h, metric) end @@ -289,15 +311,22 @@ function test_netcdf_grid_metrics_rectilinear(arch, FT) ds_s = NCDataset(filepath_sliced) @test haskey(ds_s, "time") - @test eltype(ds_s["time"]) == Float64 + @test eltype(ds_s["time"]) == Float64 # All dimensions should be Float64 by default - for var in (dims..., metrics..., vars...) + for var in (metrics..., vars...) @test haskey(ds_s, var) @test haskey(ds_s[var].attrib, "long_name") @test haskey(ds_s[var].attrib, "units") @test eltype(ds_s[var]) == FT end + for var in dims + @test haskey(ds_s, var) + @test haskey(ds_s[var].attrib, "long_name") + @test haskey(ds_s[var].attrib, "units") + @test eltype(ds_s[var]) == Float64 # All dimensions should be Float64 by default + end + @test dimsize(ds_s[:x_faa]) == (x_faa=nx,) @test dimsize(ds_s[:x_caa]) == (x_caa=nx,) @test dimsize(ds_s[:y_afa]) == (y_afa=ny,) @@ -410,7 +439,7 @@ function test_netcdf_grid_metrics_latlon(arch, FT) ds_mh = NCDataset(filepath_metrics_halos) @test haskey(ds_mh, "time") - @test eltype(ds_mh["time"]) == Float64 + @test eltype(ds_mh["time"]) == Float64 # All dimensions should be Float64 by default dims = ("λ_faa", "λ_caa", "φ_afa", "φ_aca", "z_aaf", "z_aac") metrics = ("Δλ_faa", "Δλ_caa", "Δλ_afa", "Δλ_aca", "Δz_aaf", "Δz_aac", @@ -418,13 +447,20 @@ function test_netcdf_grid_metrics_latlon(arch, FT) "Δy_ffa", "Δy_fca", "Δy_cfa", "Δy_cca") vars = ("u", "v", "w", "T", "S") - for var in (dims..., metrics..., vars...) + for var in (metrics..., vars...) @test haskey(ds_mh, var) @test haskey(ds_mh[var].attrib, "long_name") @test haskey(ds_mh[var].attrib, "units") @test eltype(ds_mh[var]) == FT end + for var in dims + @test haskey(ds_mh, var) + @test haskey(ds_mh[var].attrib, "long_name") + @test haskey(ds_mh[var].attrib, "units") + @test eltype(ds_mh[var]) == Float64 # All dimensions should be Float64 by default + end + @test dimsize(ds_mh["time"]) == (time=Nt + 1,) @test dimsize(ds_mh[:λ_faa]) == (λ_faa=Nλ + 2Hλ + 1,) @@ -464,15 +500,22 @@ function test_netcdf_grid_metrics_latlon(arch, FT) ds_h = NCDataset(filepath_nometrics) @test haskey(ds_h, "time") - @test eltype(ds_h["time"]) == Float64 + @test eltype(ds_h["time"]) == Float64 # All dimensions should be Float64 by default - for var in (dims..., vars...) + for var in vars @test haskey(ds_h, var) @test haskey(ds_h[var].attrib, "long_name") @test haskey(ds_h[var].attrib, "units") @test eltype(ds_h[var]) == FT end + for var in dims + @test haskey(ds_h, var) + @test haskey(ds_h[var].attrib, "long_name") + @test haskey(ds_h[var].attrib, "units") + @test eltype(ds_h[var]) == Float64 # All dimensions should be Float64 by default + end + # Verify that metrics are not present for metric in metrics @test !haskey(ds_h, metric) @@ -499,13 +542,20 @@ function test_netcdf_grid_metrics_latlon(arch, FT) # Test NetCDF output with metrics but no halos ds_m = NCDataset(filepath_metrics_nohalos) - for var in (dims..., metrics..., vars...) + for var in (metrics..., vars...) @test haskey(ds_m, var) @test haskey(ds_m[var].attrib, "long_name") @test haskey(ds_m[var].attrib, "units") @test eltype(ds_m[var]) == FT end + for var in dims + @test haskey(ds_m, var) + @test haskey(ds_m[var].attrib, "long_name") + @test haskey(ds_m[var].attrib, "units") + @test eltype(ds_m[var]) == Float64 # All dimensions should be Float64 by default + end + @test dimsize(ds_m[:λ_faa]) == (λ_faa=Nλ + 1,) @test dimsize(ds_m[:λ_caa]) == (λ_caa=Nλ,) @test dimsize(ds_m[:φ_afa]) == (φ_afa=Nφ + 1,) @@ -542,13 +592,20 @@ function test_netcdf_grid_metrics_latlon(arch, FT) # Test NetCDF sliced output with metrics ds_s = NCDataset(filepath_sliced) - for var in (dims..., metrics..., vars...) + for var in (metrics..., vars...) @test haskey(ds_s, var) @test haskey(ds_s[var].attrib, "long_name") @test haskey(ds_s[var].attrib, "units") @test eltype(ds_s[var]) == FT end + for var in dims + @test haskey(ds_s, var) + @test haskey(ds_s[var].attrib, "long_name") + @test haskey(ds_s[var].attrib, "units") + @test eltype(ds_s[var]) == Float64 # All dimensions should be Float64 by default + end + @test dimsize(ds_s[:λ_faa]) == (λ_faa=nx + 1,) @test dimsize(ds_s[:λ_caa]) == (λ_caa=nx,) @test dimsize(ds_s[:φ_afa]) == (φ_afa=ny,) @@ -1404,13 +1461,15 @@ function test_netcdf_rectilinear_column(arch) return nothing end -function test_thermal_bubble_netcdf_output(arch, FT) +function test_thermal_bubble_netcdf_output(arch, FT; with_halos=false) Nx, Ny, Nz = 16, 16, 16 Lx, Ly, Lz = 100, 100, 100 + Hx, Hy, Hz = 4, 3, 2 grid = RectilinearGrid(arch, topology = (Periodic, Periodic, Bounded), size = (Nx, Ny, Nz), + halo = (Hx, Hy, Hz), extent = (Lx, Ly, Lz) ) @@ -1438,13 +1497,15 @@ function test_thermal_bubble_netcdf_output(arch, FT) ) Arch = typeof(arch) - nc_filepath = "test_thermal_bubble_$(Arch)_$FT.nc" + halo_suffix = with_halos ? "_with_halos" : "" + nc_filepath = "test_thermal_bubble$(halo_suffix)_$(Arch)_$FT.nc" isfile(nc_filepath) && rm(nc_filepath) nc_writer = NetCDFWriter(model, outputs, filename = nc_filepath, schedule = IterationInterval(10), array_type = Array{FT}, + with_halos = with_halos, include_grid_metrics = false, verbose = true) @@ -1471,185 +1532,6 @@ function test_thermal_bubble_netcdf_output(arch, FT) run!(simulation) - ds3 = NCDataset(nc_filepath) - - @test haskey(ds3.attrib, "date") - @test haskey(ds3.attrib, "Julia") - @test haskey(ds3.attrib, "Oceananigans") - @test haskey(ds3.attrib, "schedule") - @test haskey(ds3.attrib, "interval") - @test haskey(ds3.attrib, "output iteration interval") - - @test !isnothing(ds3.attrib["date"]) - @test !isnothing(ds3.attrib["Julia"]) - @test !isnothing(ds3.attrib["Oceananigans"]) - @test ds3.attrib["schedule"] == "IterationInterval" - @test ds3.attrib["interval"] == 10 - @test !isnothing(ds3.attrib["output iteration interval"]) - - @test eltype(ds3["time"]) == Float64 - - @test eltype(ds3["x_caa"]) == FT - @test eltype(ds3["x_faa"]) == FT - @test eltype(ds3["y_aca"]) == FT - @test eltype(ds3["y_afa"]) == FT - @test eltype(ds3["z_aac"]) == FT - @test eltype(ds3["z_aaf"]) == FT - - @test length(ds3["x_caa"]) == Nx - @test length(ds3["y_aca"]) == Ny - @test length(ds3["z_aac"]) == Nz - @test length(ds3["x_faa"]) == Nx - @test length(ds3["y_afa"]) == Ny - @test length(ds3["z_aaf"]) == Nz+1 # z is Bounded - - @test ds3["x_caa"][1] == grid.xᶜᵃᵃ[1] - @test ds3["x_faa"][1] == grid.xᶠᵃᵃ[1] - @test ds3["y_aca"][1] == grid.yᵃᶜᵃ[1] - @test ds3["y_afa"][1] == grid.yᵃᶠᵃ[1] - @test ds3["z_aac"][1] == grid.z.cᵃᵃᶜ[1] - @test ds3["z_aaf"][1] == grid.z.cᵃᵃᶠ[1] - - @test ds3["x_caa"][end] == grid.xᶜᵃᵃ[Nx] - @test ds3["x_faa"][end] == grid.xᶠᵃᵃ[Nx] - @test ds3["y_aca"][end] == grid.yᵃᶜᵃ[Ny] - @test ds3["y_afa"][end] == grid.yᵃᶠᵃ[Ny] - @test ds3["z_aac"][end] == grid.z.cᵃᵃᶜ[Nz] - @test ds3["z_aaf"][end] == grid.z.cᵃᵃᶠ[Nz+1] # z is Bounded - - @test eltype(ds3["u"]) == FT - @test eltype(ds3["v"]) == FT - @test eltype(ds3["w"]) == FT - @test eltype(ds3["T"]) == FT - @test eltype(ds3["S"]) == FT - - u = ds3["u"][:, :, :, end] - v = ds3["v"][:, :, :, end] - w = ds3["w"][:, :, :, end] - T = ds3["T"][:, :, :, end] - S = ds3["S"][:, :, :, end] - - close(ds3) - - @test all(u .≈ Array(interior(model.velocities.u))) - @test all(v .≈ Array(interior(model.velocities.v))) - @test all(w .≈ Array(interior(model.velocities.w))) - @test all(T .≈ Array(interior(model.tracers.T))) - @test all(S .≈ Array(interior(model.tracers.S))) - - ds2 = NCDataset(nc_sliced_filepath) - - @test haskey(ds2.attrib, "date") - @test haskey(ds2.attrib, "Julia") - @test haskey(ds2.attrib, "Oceananigans") - @test haskey(ds2.attrib, "schedule") - @test haskey(ds2.attrib, "interval") - @test haskey(ds2.attrib, "output iteration interval") - - @test !isnothing(ds2.attrib["date"]) - @test !isnothing(ds2.attrib["Julia"]) - @test !isnothing(ds2.attrib["Oceananigans"]) - @test ds2.attrib["schedule"] == "IterationInterval" - @test ds2.attrib["interval"] == 10 - @test !isnothing(ds2.attrib["output iteration interval"]) - - @test eltype(ds2["time"]) == Float64 - - @test eltype(ds2["x_caa"]) == FT - @test eltype(ds2["x_faa"]) == FT - @test eltype(ds2["y_aca"]) == FT - @test eltype(ds2["y_afa"]) == FT - @test eltype(ds2["z_aac"]) == FT - @test eltype(ds2["z_aaf"]) == FT - - @test length(ds2["x_caa"]) == length(i_slice) - @test length(ds2["x_faa"]) == length(i_slice) - @test length(ds2["y_aca"]) == length(j_slice) - @test length(ds2["y_afa"]) == length(j_slice) - @test length(ds2["z_aac"]) == length(k_slice) - @test length(ds2["z_aaf"]) == length(k_slice) - - @test ds2["x_caa"][1] == grid.xᶜᵃᵃ[i_slice[1]] - @test ds2["x_faa"][1] == grid.xᶠᵃᵃ[i_slice[1]] - @test ds2["y_aca"][1] == grid.yᵃᶜᵃ[j_slice[1]] - @test ds2["y_afa"][1] == grid.yᵃᶠᵃ[j_slice[1]] - @test ds2["z_aac"][1] == grid.z.cᵃᵃᶜ[k_slice[1]] - @test ds2["z_aaf"][1] == grid.z.cᵃᵃᶠ[k_slice[1]] - - @test ds2["x_caa"][end] == grid.xᶜᵃᵃ[i_slice[end]] - @test ds2["x_faa"][end] == grid.xᶠᵃᵃ[i_slice[end]] - @test ds2["y_aca"][end] == grid.yᵃᶜᵃ[j_slice[end]] - @test ds2["y_afa"][end] == grid.yᵃᶠᵃ[j_slice[end]] - @test ds2["z_aac"][end] == grid.z.cᵃᵃᶜ[k_slice[end]] - @test ds2["z_aaf"][end] == grid.z.cᵃᵃᶠ[k_slice[end]] - - @test eltype(ds2["u"]) == FT - @test eltype(ds2["v"]) == FT - @test eltype(ds2["w"]) == FT - @test eltype(ds2["T"]) == FT - @test eltype(ds2["S"]) == FT - - u_sliced = ds2["u"][:, :, :, end] - v_sliced = ds2["v"][:, :, :, end] - w_sliced = ds2["w"][:, :, :, end] - T_sliced = ds2["T"][:, :, :, end] - S_sliced = ds2["S"][:, :, :, end] - - close(ds2) - - @test all(u_sliced .≈ Array(interior(model.velocities.u))[i_slice, j_slice, k_slice]) - @test all(v_sliced .≈ Array(interior(model.velocities.v))[i_slice, j_slice, k_slice]) - @test all(w_sliced .≈ Array(interior(model.velocities.w))[i_slice, j_slice, k_slice]) - @test all(T_sliced .≈ Array(interior(model.tracers.T))[i_slice, j_slice, k_slice]) - @test all(S_sliced .≈ Array(interior(model.tracers.S))[i_slice, j_slice, k_slice]) - - rm(nc_filepath) - rm(nc_sliced_filepath) - - return nothing -end - -function test_thermal_bubble_netcdf_output_with_halos(arch, FT) - Nx, Ny, Nz = 16, 16, 16 - Lx, Ly, Lz = 100, 100, 100 - Hx, Hy, Hz = 4, 3, 2 - - grid = RectilinearGrid(arch, - topology = (Periodic, Periodic, Bounded), - size = (Nx, Ny, Nz), - halo = (Hx, Hy, Hz), - extent = (Lx, Ly, Lz), - ) - - model = NonhydrostaticModel(; grid, - closure = ScalarDiffusivity(ν=4e-2, κ=4e-2), - buoyancy = SeawaterBuoyancy(), - tracers = (:T, :S) - ) - - simulation = Simulation(model, Δt=6, stop_iteration=10) - - # Add a cube-shaped warm temperature anomaly that takes up the middle 50% - # of the domain volume. - i1, i2 = round(Int, Nx/4), round(Int, 3Nx/4) - j1, j2 = round(Int, Ny/4), round(Int, 3Ny/4) - k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4) - view(model.tracers.T, i1:i2, j1:j2, k1:k2) .+= 0.01 - - Arch = typeof(arch) - nc_filepath = "test_thermal_bubble_with_halos_$Arch.nc" - - nc_writer = NetCDFWriter(model, - merge(model.velocities, model.tracers), - filename = nc_filepath, - schedule = IterationInterval(10), - array_type = Array{FT}, - with_halos = true) - - push!(simulation.output_writers, nc_writer) - - run!(simulation) - ds = NCDataset(nc_filepath) @test haskey(ds.attrib, "date") @@ -1666,35 +1548,58 @@ function test_thermal_bubble_netcdf_output_with_halos(arch, FT) @test ds.attrib["interval"] == 10 @test !isnothing(ds.attrib["output iteration interval"]) - @test eltype(ds["time"]) == Float64 + # Dimensions should be Float64 by default + @test eltype(ds["time"]) == Float64 + @test eltype(ds["x_caa"]) == Float64 + @test eltype(ds["x_faa"]) == Float64 + @test eltype(ds["y_aca"]) == Float64 + @test eltype(ds["y_afa"]) == Float64 + @test eltype(ds["z_aac"]) == Float64 + @test eltype(ds["z_aaf"]) == Float64 - @test eltype(ds["x_caa"]) == FT - @test eltype(ds["x_faa"]) == FT - @test eltype(ds["y_aca"]) == FT - @test eltype(ds["y_afa"]) == FT - @test eltype(ds["z_aac"]) == FT - @test eltype(ds["z_aaf"]) == FT - - @test length(ds["x_caa"]) == Nx+2Hx - @test length(ds["y_aca"]) == Ny+2Hy - @test length(ds["z_aac"]) == Nz+2Hz - @test length(ds["x_faa"]) == Nx+2Hx - @test length(ds["y_afa"]) == Ny+2Hy - @test length(ds["z_aaf"]) == Nz+2Hz+1 # z is Bounded - - @test ds["x_caa"][1] == grid.xᶜᵃᵃ[1-Hx] - @test ds["x_faa"][1] == grid.xᶠᵃᵃ[1-Hx] - @test ds["y_aca"][1] == grid.yᵃᶜᵃ[1-Hy] - @test ds["y_afa"][1] == grid.yᵃᶠᵃ[1-Hy] - @test ds["z_aac"][1] == grid.z.cᵃᵃᶜ[1-Hz] - @test ds["z_aaf"][1] == grid.z.cᵃᵃᶠ[1-Hz] - - @test ds["x_caa"][end] == grid.xᶜᵃᵃ[Nx+Hx] - @test ds["x_faa"][end] == grid.xᶠᵃᵃ[Nx+Hx] - @test ds["y_aca"][end] == grid.yᵃᶜᵃ[Ny+Hy] - @test ds["y_afa"][end] == grid.yᵃᶠᵃ[Ny+Hy] - @test ds["z_aac"][end] == grid.z.cᵃᵃᶜ[Nz+Hz] - @test ds["z_aaf"][end] == grid.z.cᵃᵃᶠ[Nz+Hz+1] # z is Bounded + if with_halos + @test length(ds["x_caa"]) == Nx+2Hx + @test length(ds["y_aca"]) == Ny+2Hy + @test length(ds["z_aac"]) == Nz+2Hz + @test length(ds["x_faa"]) == Nx+2Hx + @test length(ds["y_afa"]) == Ny+2Hy + @test length(ds["z_aaf"]) == Nz+2Hz+1 # z is Bounded + + @test ds["x_caa"][1] == grid.xᶜᵃᵃ[1-Hx] + @test ds["x_faa"][1] == grid.xᶠᵃᵃ[1-Hx] + @test ds["y_aca"][1] == grid.yᵃᶜᵃ[1-Hy] + @test ds["y_afa"][1] == grid.yᵃᶠᵃ[1-Hy] + @test ds["z_aac"][1] == grid.z.cᵃᵃᶜ[1-Hz] + @test ds["z_aaf"][1] == grid.z.cᵃᵃᶠ[1-Hz] + + @test ds["x_caa"][end] == grid.xᶜᵃᵃ[Nx+Hx] + @test ds["x_faa"][end] == grid.xᶠᵃᵃ[Nx+Hx] + @test ds["y_aca"][end] == grid.yᵃᶜᵃ[Ny+Hy] + @test ds["y_afa"][end] == grid.yᵃᶠᵃ[Ny+Hy] + @test ds["z_aac"][end] == grid.z.cᵃᵃᶜ[Nz+Hz] + @test ds["z_aaf"][end] == grid.z.cᵃᵃᶠ[Nz+Hz+1] # z is Bounded + else + @test length(ds["x_caa"]) == Nx + @test length(ds["y_aca"]) == Ny + @test length(ds["z_aac"]) == Nz + @test length(ds["x_faa"]) == Nx + @test length(ds["y_afa"]) == Ny + @test length(ds["z_aaf"]) == Nz+1 # z is Bounded + + @test ds["x_caa"][1] == grid.xᶜᵃᵃ[1] + @test ds["x_faa"][1] == grid.xᶠᵃᵃ[1] + @test ds["y_aca"][1] == grid.yᵃᶜᵃ[1] + @test ds["y_afa"][1] == grid.yᵃᶠᵃ[1] + @test ds["z_aac"][1] == grid.z.cᵃᵃᶜ[1] + @test ds["z_aaf"][1] == grid.z.cᵃᵃᶠ[1] + + @test ds["x_caa"][end] == grid.xᶜᵃᵃ[Nx] + @test ds["x_faa"][end] == grid.xᶠᵃᵃ[Nx] + @test ds["y_aca"][end] == grid.yᵃᶜᵃ[Ny] + @test ds["y_afa"][end] == grid.yᵃᶠᵃ[Ny] + @test ds["z_aac"][end] == grid.z.cᵃᵃᶜ[Nz] + @test ds["z_aaf"][end] == grid.z.cᵃᵃᶠ[Nz+1] # z is Bounded + end @test eltype(ds["u"]) == FT @test eltype(ds["v"]) == FT @@ -1710,14 +1615,93 @@ function test_thermal_bubble_netcdf_output_with_halos(arch, FT) close(ds) - @test all(u .≈ Array(model.velocities.u.data.parent)) - @test all(v .≈ Array(model.velocities.v.data.parent)) - @test all(w .≈ Array(model.velocities.w.data.parent)) - @test all(T .≈ Array(model.tracers.T.data.parent)) - @test all(S .≈ Array(model.tracers.S.data.parent)) + if with_halos + @test all(u .≈ Array(model.velocities.u.data.parent)) + @test all(v .≈ Array(model.velocities.v.data.parent)) + @test all(w .≈ Array(model.velocities.w.data.parent)) + @test all(T .≈ Array(model.tracers.T.data.parent)) + @test all(S .≈ Array(model.tracers.S.data.parent)) + else + @test all(u .≈ Array(interior(model.velocities.u))) + @test all(v .≈ Array(interior(model.velocities.v))) + @test all(w .≈ Array(interior(model.velocities.w))) + @test all(T .≈ Array(interior(model.tracers.T))) + @test all(S .≈ Array(interior(model.tracers.S))) + end rm(nc_filepath) + # Test sliced output only when not using halos + if !with_halos + ds2 = NCDataset(nc_sliced_filepath) + + @test haskey(ds2.attrib, "date") + @test haskey(ds2.attrib, "Julia") + @test haskey(ds2.attrib, "Oceananigans") + @test haskey(ds2.attrib, "schedule") + @test haskey(ds2.attrib, "interval") + @test haskey(ds2.attrib, "output iteration interval") + + @test !isnothing(ds2.attrib["date"]) + @test !isnothing(ds2.attrib["Julia"]) + @test !isnothing(ds2.attrib["Oceananigans"]) + @test ds2.attrib["schedule"] == "IterationInterval" + @test ds2.attrib["interval"] == 10 + @test !isnothing(ds2.attrib["output iteration interval"]) + + # Dimensions should be Float64 by default + @test eltype(ds2["time"]) == Float64 + @test eltype(ds2["x_caa"]) == Float64 + @test eltype(ds2["x_faa"]) == Float64 + @test eltype(ds2["y_aca"]) == Float64 + @test eltype(ds2["y_afa"]) == Float64 + @test eltype(ds2["z_aac"]) == Float64 + @test eltype(ds2["z_aaf"]) == Float64 + + @test length(ds2["x_caa"]) == length(i_slice) + @test length(ds2["x_faa"]) == length(i_slice) + @test length(ds2["y_aca"]) == length(j_slice) + @test length(ds2["y_afa"]) == length(j_slice) + @test length(ds2["z_aac"]) == length(k_slice) + @test length(ds2["z_aaf"]) == length(k_slice) + + @test ds2["x_caa"][1] == grid.xᶜᵃᵃ[i_slice[1]] + @test ds2["x_faa"][1] == grid.xᶠᵃᵃ[i_slice[1]] + @test ds2["y_aca"][1] == grid.yᵃᶜᵃ[j_slice[1]] + @test ds2["y_afa"][1] == grid.yᵃᶠᵃ[j_slice[1]] + @test ds2["z_aac"][1] == grid.z.cᵃᵃᶜ[k_slice[1]] + @test ds2["z_aaf"][1] == grid.z.cᵃᵃᶠ[k_slice[1]] + + @test ds2["x_caa"][end] == grid.xᶜᵃᵃ[i_slice[end]] + @test ds2["x_faa"][end] == grid.xᶠᵃᵃ[i_slice[end]] + @test ds2["y_aca"][end] == grid.yᵃᶜᵃ[j_slice[end]] + @test ds2["y_afa"][end] == grid.yᵃᶠᵃ[j_slice[end]] + @test ds2["z_aac"][end] == grid.z.cᵃᵃᶜ[k_slice[end]] + @test ds2["z_aaf"][end] == grid.z.cᵃᵃᶠ[k_slice[end]] + + @test eltype(ds2["u"]) == FT + @test eltype(ds2["v"]) == FT + @test eltype(ds2["w"]) == FT + @test eltype(ds2["T"]) == FT + @test eltype(ds2["S"]) == FT + + u_sliced = ds2["u"][:, :, :, end] + v_sliced = ds2["v"][:, :, :, end] + w_sliced = ds2["w"][:, :, :, end] + T_sliced = ds2["T"][:, :, :, end] + S_sliced = ds2["S"][:, :, :, end] + + close(ds2) + + @test all(u_sliced .≈ Array(interior(model.velocities.u))[i_slice, j_slice, k_slice]) + @test all(v_sliced .≈ Array(interior(model.velocities.v))[i_slice, j_slice, k_slice]) + @test all(w_sliced .≈ Array(interior(model.velocities.w))[i_slice, j_slice, k_slice]) + @test all(T_sliced .≈ Array(interior(model.tracers.T))[i_slice, j_slice, k_slice]) + @test all(S_sliced .≈ Array(interior(model.tracers.S))[i_slice, j_slice, k_slice]) + + rm(nc_sliced_filepath) + end + return nothing end @@ -1891,7 +1875,8 @@ function test_netcdf_function_output(arch) dimensions = dims, array_type = Array{Float64}, include_grid_metrics = false, - verbose = true) + verbose = true, + overwrite_existing = true) run!(simulation) @@ -2121,14 +2106,14 @@ function test_netcdf_time_averaging(arch) horizontal_average_nc_filepath = "decay_averaged_field_test_$Arch.nc" simulation.output_writers[:horizontal_average] = - NetCDFWriter( - model, - nc_outputs, - array_type = Array{Float64}, - verbose = true, - filename = horizontal_average_nc_filepath, - schedule = TimeInterval(10Δt), - include_grid_metrics = false) + NetCDFWriter(model, + nc_outputs, + array_type = Array{Float64}, + verbose = true, + filename = horizontal_average_nc_filepath, + schedule = TimeInterval(10Δt), + include_grid_metrics = false, + overwrite_existing = true) multiple_time_average_nc_filepath = "decay_windowed_time_average_test_$Arch.nc" single_time_average_nc_filepath = "single_decay_windowed_time_average_test_$Arch.nc" @@ -2137,24 +2122,24 @@ function test_netcdf_time_averaging(arch) single_nc_output = Dict("c1" => ∫c1_dxdy) simulation.output_writers[:single_output_time_average] = - NetCDFWriter( - model, - single_nc_output, - array_type = Array{Float64}, - verbose = true, - filename = single_time_average_nc_filepath, - schedule = AveragedTimeInterval(10Δt; window, stride), - include_grid_metrics = false) + NetCDFWriter(model, + single_nc_output, + array_type = Array{Float64}, + verbose = true, + filename = single_time_average_nc_filepath, + schedule = AveragedTimeInterval(10Δt; window, stride), + include_grid_metrics = false, + overwrite_existing = true) simulation.output_writers[:multiple_output_time_average] = - NetCDFWriter( - model, - nc_outputs, - array_type = Array{Float64}, - verbose = true, - filename = multiple_time_average_nc_filepath, - schedule = AveragedTimeInterval(10Δt; window, stride), - include_grid_metrics = false) + NetCDFWriter(model, + nc_outputs, + array_type = Array{Float64}, + verbose = true, + filename = multiple_time_average_nc_filepath, + schedule = AveragedTimeInterval(10Δt; window, stride), + include_grid_metrics = false, + overwrite_existing = true) run!(simulation) @@ -2578,7 +2563,7 @@ function test_netcdf_free_surface_only_output(arch) ds_h = NCDataset(filepath_with_halos) @test haskey(ds_h, "η") - @test dimsize(ds_h["η"]) == (λ_caa=Nλ + 2Hλ, φ_aca=Nφ + 2Hφ, time=Nt + 1) + @test dimsize(ds_h["η"]) == (λ_caa=Nλ + 2Hλ, φ_aca=Nφ + 2Hφ, z_aaf_η=1, time=Nt + 1) close(ds_h) rm(filepath_with_halos) @@ -2586,7 +2571,7 @@ function test_netcdf_free_surface_only_output(arch) ds_n = NCDataset(filepath_no_halos) @test haskey(ds_n, "η") - @test dimsize(ds_n["η"]) == (λ_caa=Nλ, φ_aca=Nφ, time=Nt + 1) + @test dimsize(ds_n["η"]) == (λ_caa=Nλ, φ_aca=Nφ, z_aaf_η=1, time=Nt + 1) close(ds_n) rm(filepath_no_halos) @@ -2630,7 +2615,8 @@ function test_netcdf_free_surface_mixed_output(arch) NetCDFWriter(model, outputs; filename = filepath_with_halos, schedule = IterationInterval(1), - with_halos = true) + with_halos = true, + overwrite_existing = true) filepath_no_halos = "test_mixed_free_surface_no_halos_$Arch.nc" isfile(filepath_no_halos) && rm(filepath_no_halos) @@ -2639,14 +2625,15 @@ function test_netcdf_free_surface_mixed_output(arch) NetCDFWriter(model, outputs; filename = filepath_no_halos, schedule = IterationInterval(1), - with_halos = false) + with_halos = false, + overwrite_existing = true) run!(simulation) ds_h = NCDataset(filepath_with_halos) @test haskey(ds_h, "η") - @test dimsize(ds_h["η"]) == (λ_caa=Nλ + 2Hλ, φ_aca=Nφ + 2Hφ, time=Nt + 1) + @test dimsize(ds_h["η"]) == (λ_caa=Nλ + 2Hλ, φ_aca=Nφ + 2Hφ, z_aaf_η=1, time=Nt + 1) @test dimsize(ds_h[:u]) == (λ_faa=Nλ + 2Hλ + 1, φ_aca=Nφ + 2Hφ, z_aac=Nz + 2Hz, time=Nt + 1) @test dimsize(ds_h[:v]) == (λ_caa=Nλ + 2Hλ, φ_afa=Nφ + 2Hφ + 1, z_aac=Nz + 2Hz, time=Nt + 1) @@ -2660,7 +2647,7 @@ function test_netcdf_free_surface_mixed_output(arch) ds_n = NCDataset(filepath_no_halos) @test haskey(ds_n, "η") - @test dimsize(ds_n["η"]) == (λ_caa=Nλ, φ_aca=Nφ, time=Nt + 1) + @test dimsize(ds_n["η"]) == (λ_caa=Nλ, φ_aca=Nφ, z_aaf_η=1, time=Nt + 1) @test dimsize(ds_n[:u]) == (λ_faa=Nλ + 1, φ_aca=Nφ, z_aac=Nz, time=Nt + 1) @test dimsize(ds_n[:v]) == (λ_caa=Nλ, φ_afa=Nφ + 1, z_aac=Nz, time=Nt + 1) @@ -2880,6 +2867,175 @@ function test_netcdf_multiple_grids_defvar(grid1, grid2; immersed=false) return nothing end +function test_netcdf_writer_different_grid(arch) + + grid = RectilinearGrid(arch, size=(1, 1, 8), extent=(1, 1, 1)) + model = NonhydrostaticModel(; grid) + + coarse_grid = RectilinearGrid(arch, size=(grid.Nx, grid.Ny, grid.Nz÷2), extent=(grid.Lx, grid.Ly, grid.Lz)) + coarse_u = Field{Face, Center, Center}(coarse_grid) + + interpolate_u(model) = interpolate!(coarse_u, model.velocities.u) + outputs = (; u = interpolate_u) + + Arch = typeof(arch) + filepath = "test_coarse_u_$Arch.nc" + isfile(filepath) && rm(filepath) + + # This should work: NetCDFWriter should use coarse_grid for dimensions + # when grid parameter is provided, not model.grid + output_writer = NetCDFWriter(model, outputs; + grid = coarse_grid, + filename = filepath, + schedule = IterationInterval(1), + overwrite_existing = true) + + # Run simulation to write output + simulation = Simulation(model, Δt=0.1, stop_iteration=2) + simulation.output_writers[:coarse] = output_writer + run!(simulation) + + # Verify that dimensions match the coarse grid, not the model grid + ds = NCDataset(filepath) + @test length(ds["z_aac"]) == grid.Nz÷2 # Should be 4, not 8 + @test length(ds["z_aaf"]) == grid.Nz÷2 + 1 # Should be 5, not 9 + + # Verify that the z coordinates match the coarse grid + expected_z_centers = znodes(coarse_grid, Center()) + expected_z_faces = znodes(coarse_grid, Face()) + @test ds["z_aac"][:] ≈ Array(expected_z_centers) + @test ds["z_aaf"][:] ≈ Array(expected_z_faces) + + close(ds) + rm(filepath) + + return nothing +end + +function test_singleton_dimension_behavior(arch) + Nx, Nz = 8, 8 + Hx, Hz = 2, 3 + Lx, H = 2, 1 + + # Create grid with Flat in y dimension + grid = RectilinearGrid(arch, + topology = (Periodic, Flat, Bounded), + size = (Nx, Nz), + halo = (Hx, Hz), + x = (-Lx, Lx), + z = (-H, 0)) + + model = NonhydrostaticModel(; grid, + closure = ScalarDiffusivity(ν=4e-2, κ=4e-2), + buoyancy = SeawaterBuoyancy(), + tracers = (:T, :S)) + u_xavg = Field(Average(model.velocities.u, dims=(1))) + + Nt = 5 + simulation = Simulation(model, Δt=0.1, stop_iteration=Nt) + + Arch = typeof(arch) + filepath_full = "test_singleton_full_u_$Arch.nc" + simulation.output_writers[:full] = NetCDFWriter(model, (; u = model.velocities.u, u_xavg), + filename = filepath_full, + schedule = IterationInterval(1), + array_type = Array{Float64}, + with_halos = false, + include_grid_metrics = false, + overwrite_existing = true) + + filepath_slice = "test_singleton_slice_u_$Arch.nc" + simulation.output_writers[:slice] = NetCDFWriter(model, (; u = model.velocities.u, u_xavg), + filename = filepath_slice, + indices = (:, :, 1), + schedule = IterationInterval(1), + array_type = Array{Float64}, + with_halos = false, + include_grid_metrics = false, + overwrite_existing = true) + + run!(simulation) + + # Verify full u output + ds_full = NCDataset(filepath_full) + @test haskey(ds_full, "u") + @test haskey(ds_full, "u_xavg") + @test dimsize(ds_full[:u]) == (x_faa=Nx, z_aac=Nz, time=Nt + 1) + @test dimsize(ds_full[:u_xavg]) == (z_aac=Nz, time=Nt + 1) + @test !haskey(ds_full, "y_aca") # y dimension should not exist (Flat) + @test !haskey(ds_full, "y_afa") # y dimension should not exist (Flat) + close(ds_full) + rm(filepath_full) + + # Verify sliced u output (:, :, 1) + ds_slice = NCDataset(filepath_slice) + @test haskey(ds_slice, "u") + @test haskey(ds_slice, "u_xavg") + @test dimsize(ds_slice[:u]) == (x_faa=Nx, z_aac=1, time=Nt + 1) + @test dimsize(ds_slice[:u_xavg]) == (z_aac=1, time=Nt + 1) + @test !haskey(ds_slice, "y_aca") # y dimension should not exist (Flat) + @test !haskey(ds_slice, "y_afa") # y dimension should not exist (Flat) + close(ds_slice) + rm(filepath_slice) + + return nothing +end + +function test_netcdf_dimension_type(arch) + grid = RectilinearGrid(arch, size=(4, 4, 4), extent=(1, 1, 1)) + model = NonhydrostaticModel(; grid) + + Nt = 3 + simulation = Simulation(model, Δt=0.1, stop_iteration=Nt) + + Arch = typeof(arch) + + # Test with default dimension_type (Float64) + filepath_float64 = "test_dimension_type_float64_$Arch.nc" + simulation.output_writers[:float64] = NetCDFWriter(model, (; u = model.velocities.u), + filename = filepath_float64, + schedule = IterationInterval(1), + include_grid_metrics = false, + overwrite_existing = true) + + # Test with Float32 dimension_type + filepath_float32 = "test_dimension_type_float32_$Arch.nc" + simulation.output_writers[:float32] = NetCDFWriter(model, (; u = model.velocities.u), + filename = filepath_float32, + schedule = IterationInterval(1), + dimension_type = Float32, + include_grid_metrics = false, + overwrite_existing = true) + + run!(simulation) + + # Verify Float64 dimensions (default) + ds_float64 = NCDataset(filepath_float64) + @test eltype(ds_float64["time"]) == Float64 + @test eltype(ds_float64["x_faa"]) == Float64 + @test eltype(ds_float64["x_caa"]) == Float64 + @test eltype(ds_float64["y_aca"]) == Float64 + @test eltype(ds_float64["y_afa"]) == Float64 + @test eltype(ds_float64["z_aac"]) == Float64 + @test eltype(ds_float64["z_aaf"]) == Float64 + close(ds_float64) + rm(filepath_float64) + + # Verify Float32 dimensions + ds_float32 = NCDataset(filepath_float32) + @test eltype(ds_float32["time"]) == Float32 + @test eltype(ds_float32["x_faa"]) == Float32 + @test eltype(ds_float32["x_caa"]) == Float32 + @test eltype(ds_float32["y_aca"]) == Float32 + @test eltype(ds_float32["y_afa"]) == Float32 + @test eltype(ds_float32["z_aac"]) == Float32 + @test eltype(ds_float32["z_aaf"]) == Float32 + close(ds_float32) + rm(filepath_float32) + + return nothing +end + for arch in archs @testset "NetCDF output writer [$(typeof(arch))]" begin @info " Testing NetCDF output writer [$(typeof(arch))]..." @@ -2895,6 +3051,7 @@ for arch in archs test_datetime_netcdf_output(arch) test_timedate_netcdf_output(arch) + test_netcdf_dimension_type(arch) test_netcdf_grid_metrics_rectilinear(arch, Float64) test_netcdf_grid_metrics_rectilinear(arch, Float32) @@ -2915,8 +3072,8 @@ for arch in archs test_thermal_bubble_netcdf_output(arch, Float64) test_thermal_bubble_netcdf_output(arch, Float32) - test_thermal_bubble_netcdf_output_with_halos(arch, Float64) - test_thermal_bubble_netcdf_output_with_halos(arch, Float32) + test_thermal_bubble_netcdf_output(arch, Float64, with_halos=true) + test_thermal_bubble_netcdf_output(arch, Float32, with_halos=true) test_netcdf_size_file_splitting(arch) test_netcdf_time_file_splitting(arch) @@ -2939,6 +3096,10 @@ for arch in archs test_netcdf_buoyancy_force(arch) + test_netcdf_writer_different_grid(arch) + + test_singleton_dimension_behavior(arch) + for grids in ((rectilinear_grid1, rectilinear_grid2), (latlon_grid1, latlon_grid2)) grid1, grid2 = grids diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl index 2e6a7055f9..576d0c37af 100644 --- a/test/test_output_readers.jl +++ b/test/test_output_readers.jl @@ -105,7 +105,7 @@ function test_pickup_with_inaccurate_times() grid = RectilinearGrid(size=(2, 2, 2), extent=(1, 1, 1)) times = collect(0:0.1:3) filename = "fts_inaccurate_times_test.jld2" - f_tmp = Field{Center,Center,Center}(grid) + f_tmp = Field{Center,Center,Center}(grid) f = FieldTimeSeries{Center, Center, Center}(grid, times; backend=OnDisk(), path=filename, name="f") for (it, time) in enumerate(f.times) @@ -313,7 +313,7 @@ end @info " Testing FieldTimeSeries pickup..." Random.seed!(1234) for n in -4:4 - Δt = (1.1 + rand()) * 10.0^n + Δt = (1.1 + rand()) * 10.0^n Lx = 10 * Δt for FT in (Float32, Float64) filename = generate_nonzero_simulation_data(Lx, Δt, FT) @@ -424,7 +424,7 @@ end end end end - + @testset "Test chunked abstraction" begin @info " Testing Chunked abstraction..." filepath = "testfile.jld2" @@ -475,7 +475,7 @@ end end end end - + for Backend in [InMemory, OnDisk] @testset "FieldDataset{$Backend} indexing" begin @info " Testing FieldDataset{$Backend} indexing..." diff --git a/test/test_output_writers.jl b/test/test_output_writers.jl index c5c1158c38..c14d727714 100644 --- a/test/test_output_writers.jl +++ b/test/test_output_writers.jl @@ -100,7 +100,6 @@ function test_creating_and_appending(model, output_writer) output = fields(model) filename = "test_creating_and_appending" - # Create a simulation with `overwrite_existing = true` and run it simulation.output_writers[:writer] = writer = output_writer(model, output, filename = filename, schedule = IterationInterval(1),