diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl index 156c3b17cb..7030e5dd59 100644 --- a/ext/OceananigansNCDatasetsExt.jl +++ b/ext/OceananigansNCDatasetsExt.jl @@ -742,27 +742,64 @@ materialize_from_netcdf(x::Number) = x materialize_from_netcdf(x::Array) = Tuple(x) materialize_from_netcdf(x::String) = @eval $(Meta.parse(x)) + +function netcdf_grid_constructor_info(grid) + underlying_grid_args, underlying_grid_kwargs = constructor_arguments(grid) + + immersed_grid_args = Dict() + + underlying_grid_type = typeof(grid).name.wrapper |> string # Save type of grid for reconstruction + grid_metadata = Dict(:immersed_boundary_type => nothing, + :underlying_grid_type => underlying_grid_type) + return underlying_grid_args, underlying_grid_kwargs, immersed_grid_args, grid_metadata +end + +function netcdf_grid_constructor_info(grid::ImmersedBoundaryGrid) + underlying_grid_args, underlying_grid_kwargs, immersed_grid_args = constructor_arguments(grid) + + immersed_boundary_type = typeof(grid.immersed_boundary).name.wrapper |> string # Save type of immersed boundary for reconstruction + underlying_grid_type = typeof(grid.underlying_grid).name.wrapper |> string # Save type of underlyinh grid for reconstruction + + grid_metadata = Dict(:immersed_boundary_type => immersed_boundary_type, + :underlying_grid_type => underlying_grid_type) + return underlying_grid_args, underlying_grid_kwargs, immersed_grid_args, grid_metadata +end + +function write_immersed_boundary_data!(ds, grid::ImmersedBoundaryGrid, immersed_grid_args) + ibg_group = defGroup(ds, "immersed_grid_reconstruction_args") + + if (grid.immersed_boundary isa GridFittedBottom) || (grid.immersed_boundary isa PartialCellBottom) + bottom_height = pop!(immersed_grid_args, :bottom_height) + defVar(ibg_group, "bottom_height", bottom_height; attrib=convert_for_netcdf(immersed_grid_args)) + + elseif grid.immersed_boundary isa GridFittedBoundary + mask = pop!(immersed_grid_args, :mask) + defVar(ibg_group, "mask", mask; attrib=convert_for_netcdf(immersed_grid_args)) + end + return ds +end + +write_immersed_boundary_data!(ds, grid, immersed_grid_args) = nothing + function write_grid_reconstruction_data!(ds, grid; array_type=Array{eltype(grid)}, deflatelevel=0) - grid_attrs, grid_dims = grid_attributes(grid) + # Save buman-readable grid attributes for inspection. Not used for reconstruction. + # TODO: Remove this once we have a better way to save grid attributes. + grid_attrs, grid_dims = grid_attributes(grid) sorted_grid_attrs = sort(collect(pairs(grid_attrs)), by=first) # Organizes attributes to make it more easily human-readable ds_grid = defGroup(ds, "grid_attributes"; attrib = sorted_grid_attrs) for (dim_name, dim_array) in grid_dims defVar(ds_grid, dim_name, array_type(dim_array), (dim_name,); deflatelevel) end - args, kwargs = constructor_arguments(grid) + underlying_grid_args, underlying_grid_kwargs, immersed_grid_args, grid_metadata = netcdf_grid_constructor_info(grid) + underlying_grid_args, underlying_grid_kwargs, grid_metadata = map(convert_for_netcdf, (underlying_grid_args, underlying_grid_kwargs, grid_metadata)) - # This is needed for now to prevent NetCDF from throwing an error, but precludes - # us from reconstructing immersed grids from NetCDF. This is a temporary fix. - :mask ∈ keys(args) && delete!(args, :mask) - :bottom_height ∈ keys(args) && delete!(args, :bottom_height) + defGroup(ds, "underlying_grid_reconstruction_args"; attrib = underlying_grid_args) + defGroup(ds, "underlying_grid_reconstruction_kwargs"; attrib = underlying_grid_kwargs) + defGroup(ds, "grid_reconstruction_metadata"; attrib = grid_metadata) - args, kwargs = map(convert_for_netcdf, (args, kwargs)) - args["grid_type"] = typeof(grid).name.wrapper |> string # Save type of grid for reconstruction - - defGroup(ds, "grid_reconstruction_args"; attrib = args) - defGroup(ds, "grid_reconstruction_kwargs"; attrib = kwargs) + write_immersed_boundary_data!(ds, grid, immersed_grid_args) return ds end @@ -774,23 +811,50 @@ function reconstruct_grid(filename::String) return grid end +function reconstruct_immersed_boundary(ds) + ibg_group = ds.group["immersed_grid_reconstruction_args"] + + grid_reconstruction_metadata = ds.group["grid_reconstruction_metadata"].attrib |> materialize_from_netcdf + immersed_boundary_type = grid_reconstruction_metadata[:immersed_boundary_type] + if immersed_boundary_type == GridFittedBottom + bottom_height = Array(ibg_group["bottom_height"]) + immersed_boundary = immersed_boundary_type(bottom_height) + # immersed_grid_reconstruction_kwargs = ibg_group.attrib |> materialize_from_netcdf + # immersed_boundary = immersed_boundary_type(bottom_height; immersed_grid_reconstruction_kwargs...) + + elseif immersed_boundary_type == PartialCellBottom + bottom_height = Array(ibg_group["bottom_height"]) + immersed_boundary = immersed_boundary_type(bottom_height) + + elseif immersed_boundary_type == GridFittedBoundary + mask = Array(ibg_group["mask"]) + immersed_boundary = immersed_boundary_type(mask) + + else + error("Unsupported immersed boundary type: $immersed_boundary_type") + end + return immersed_boundary +end + + function reconstruct_grid(ds) # Read back the grid reconstruction metadata - grid_reconstruction_args = ds.group["grid_reconstruction_args"].attrib |> materialize_from_netcdf - grid_reconstruction_kwargs = ds.group["grid_reconstruction_kwargs"].attrib |> materialize_from_netcdf - - # Pop out infomration about grid type and immersed boundary type from Dict - grid_type = pop!(grid_reconstruction_args, :grid_type) - is_immersed = haskey(grid_reconstruction_args, :immersed_boundary_type) - if is_immersed - ib_type = pop!(grid_reconstruction_args, :immersed_boundary_type) - end + underlying_grid_reconstruction_args = ds.group["underlying_grid_reconstruction_args"].attrib |> materialize_from_netcdf + underlying_grid_reconstruction_kwargs = ds.group["underlying_grid_reconstruction_kwargs"].attrib |> materialize_from_netcdf + grid_reconstruction_metadata = ds.group["grid_reconstruction_metadata"].attrib |> materialize_from_netcdf + + # Pop out infomration about the underlying grid + underlying_grid_type = grid_reconstruction_metadata[:underlying_grid_type] + underlying_grid = underlying_grid_type(values(underlying_grid_reconstruction_args)...; underlying_grid_reconstruction_kwargs...) - # Reconstruct the grid which may or may not be an underlying grid to an ImmersedBoundaryGrid - maybe_underlying_grid = grid_type(values(grid_reconstruction_args)...; grid_reconstruction_kwargs...) + # If this is an ImmersedBoundaryGrid, reconstruct the immersed boundary, otherwise underlying grid is the final grid + if grid_reconstruction_metadata[:immersed_boundary_type] isa Nothing + grid = underlying_grid + else + immersed_boundary = reconstruct_immersed_boundary(ds) + grid = ImmersedBoundaryGrid(underlying_grid, immersed_boundary) + end - # If this is an ImmersedBoundaryGrid, reconstruct the immersed boundary - grid = is_immersed ? ImmersedBoundaryGrid(maybe_underlying_grid, ib_type) : maybe_underlying_grid return grid end diff --git a/src/Grids/abstract_grid.jl b/src/Grids/abstract_grid.jl index d8c0ec4b40..fc41ece9ad 100644 --- a/src/Grids/abstract_grid.jl +++ b/src/Grids/abstract_grid.jl @@ -73,7 +73,7 @@ Return the architecture that the `grid` lives on. Return a 3-tuple of the number of "center" cells on a grid in (x, y, z). Center cells have the location (Center, Center, Center). """ -@inline Base.size(grid::AbstractGrid) = (grid.Nx, grid.Ny, grid.Nz) +@inline Base.size(grid::AbstractGrid) = map(Int, (grid.Nx, grid.Ny, grid.Nz)) Base.eltype(::AbstractGrid{FT}) where FT = FT Base.eltype(::Type{<:Oceananigans.Grids.AbstractGrid{FT}}) where FT = FT Base.eps(::AbstractGrid{FT}) where FT = eps(FT) @@ -96,7 +96,7 @@ end Return a 3-tuple with the number of halo cells on either side of the domain in (x, y, z). """ -halo_size(grid) = (grid.Hx, grid.Hy, grid.Hz) +halo_size(grid) = map(Int, (grid.Hx, grid.Hy, grid.Hz)) halo_size(grid, d) = halo_size(grid)[d] @inline Base.size(grid::AbstractGrid, d::Int) = size(grid)[d] diff --git a/src/ImmersedBoundaries/grid_fitted_bottom.jl b/src/ImmersedBoundaries/grid_fitted_bottom.jl index 3367c3a9b0..e1afc6ed68 100644 --- a/src/ImmersedBoundaries/grid_fitted_bottom.jl +++ b/src/ImmersedBoundaries/grid_fitted_bottom.jl @@ -142,6 +142,7 @@ end ##### Static column depth ##### +# AbstractGridFittedBottomImmersedBoundaryGrid const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom} @inline static_column_depthᶜᶜᵃ(i, j, ibg::AGFBIBG) = @inbounds rnode(i, j, ibg.Nz+1, ibg, c, c, f) - ibg.immersed_boundary.bottom_height[i, j, 1] @@ -160,11 +161,10 @@ YFlatAGFIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Flat, <:Any, <:Any, <:Abstrac function constructor_arguments(grid::AGFBIBG) - args, kwargs = constructor_arguments(grid.underlying_grid) - args = merge(args, Dict(:bottom_height => grid.immersed_boundary.bottom_height, - :immersed_condition => grid.immersed_boundary.immersed_condition, - :immersed_boundary_type => nameof(typeof(grid.immersed_boundary)))) - return args, kwargs + underlying_grid_args, underlying_grid_kwargs = constructor_arguments(grid.underlying_grid) + grid_fitted_bottom_args = Dict(:bottom_height => grid.immersed_boundary.bottom_height, + :immersed_condition => grid.immersed_boundary.immersed_condition) + return underlying_grid_args, underlying_grid_kwargs, grid_fitted_bottom_args end function Base.:(==)(gfb1::GridFittedBottom, gfb2::GridFittedBottom) diff --git a/src/ImmersedBoundaries/grid_fitted_boundary.jl b/src/ImmersedBoundaries/grid_fitted_boundary.jl index 03d061891e..28cbf2e27a 100644 --- a/src/ImmersedBoundaries/grid_fitted_boundary.jl +++ b/src/ImmersedBoundaries/grid_fitted_boundary.jl @@ -37,10 +37,9 @@ Adapt.adapt_structure(to, ib::AbstractGridFittedBoundary) = GridFittedBoundary(a const AGFBoundIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:GridFittedBoundary} function constructor_arguments(grid::AGFBoundIBG) - args, kwargs = constructor_arguments(grid.underlying_grid) - args = merge(args, Dict(:mask => grid.immersed_boundary.mask, - :immersed_boundary_type => nameof(typeof(grid.immersed_boundary)))) - return args, kwargs + underlying_grid_args, underlying_grid_kwargs = constructor_arguments(grid.underlying_grid) + grid_fitted_boundary_args = Dict(:mask => grid.immersed_boundary.mask) + return underlying_grid_args, underlying_grid_kwargs, grid_fitted_boundary_args end Base.:(==)(gfb1::GridFittedBoundary, gfb2::GridFittedBoundary) = gfb1.mask == gfb2.mask \ No newline at end of file diff --git a/src/ImmersedBoundaries/partial_cell_bottom.jl b/src/ImmersedBoundaries/partial_cell_bottom.jl index e51694c4fe..70fd826f30 100644 --- a/src/ImmersedBoundaries/partial_cell_bottom.jl +++ b/src/ImmersedBoundaries/partial_cell_bottom.jl @@ -210,11 +210,10 @@ VSPCBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:AbstractStaticGrid @inline Δzᶠᶠᶠ(i, j, k, ibg::VSPCBIBG) = Δrᶠᶠᶠ(i, j, k, ibg) function constructor_arguments(grid::PCBIBG) - args, kwargs = constructor_arguments(grid.underlying_grid) - args = merge(args, Dict(:bottom_height => grid.immersed_boundary.bottom_height, - :minimum_fractional_cell_height => grid.immersed_boundary.minimum_fractional_cell_height, - :immersed_boundary_type => nameof(typeof(grid.immersed_boundary)))) - return args, kwargs + underlying_grid_args, underlying_grid_kwargs = constructor_arguments(grid.underlying_grid) + partial_cell_bottom_args = Dict(:bottom_height => grid.immersed_boundary.bottom_height, + :minimum_fractional_cell_height => grid.immersed_boundary.minimum_fractional_cell_height) + return underlying_grid_args, underlying_grid_kwargs, partial_cell_bottom_args end function Base.:(==)(pcb1::PartialCellBottom, pcb2::PartialCellBottom) diff --git a/test/test_grid_reconstruction.jl b/test/test_grid_reconstruction.jl index 39f4713570..2405a0a179 100644 --- a/test/test_grid_reconstruction.jl +++ b/test/test_grid_reconstruction.jl @@ -227,22 +227,28 @@ function test_latitude_longitude_grid_reconstruction(original_grid) end function test_immersed_grid_reconstruction(original_grid) - args, kwargs = constructor_arguments(original_grid) + underlying_grid_args, underlying_grid_kwargs, immersed_boundary_args = constructor_arguments(original_grid) - @test :immersed_boundary_type in keys(args) - @test :architecture in keys(args) - @test :number_type in keys(args) + @test :architecture in keys(underlying_grid_args) + @test :number_type in keys(underlying_grid_args) # Reconstruct the immersed boundary and then the grid original_ib = original_grid.immersed_boundary if original_ib isa GridFittedBottom - reconstructed_ib = GridFittedBottom(args[:bottom_height], args[:immersed_condition]) + @test :bottom_height in keys(immersed_boundary_args) + @test :immersed_condition in keys(immersed_boundary_args) + reconstructed_ib = GridFittedBottom(immersed_boundary_args[:bottom_height], immersed_boundary_args[:immersed_condition]) + elseif original_ib isa PartialCellBottom - reconstructed_ib = PartialCellBottom(args[:bottom_height], args[:minimum_fractional_cell_height]) + @test :bottom_height in keys(immersed_boundary_args) + @test :minimum_fractional_cell_height in keys(immersed_boundary_args) + reconstructed_ib = PartialCellBottom(immersed_boundary_args[:bottom_height], immersed_boundary_args[:minimum_fractional_cell_height]) + elseif original_ib isa GridFittedBoundary - reconstructed_ib = GridFittedBoundary(args[:mask]) + @test :mask in keys(immersed_boundary_args) + reconstructed_ib = GridFittedBoundary(immersed_boundary_args[:mask]) end - reconstructed_underlying_grid = RectilinearGrid(args[:architecture], args[:number_type]; kwargs...) + reconstructed_underlying_grid = RectilinearGrid(values(underlying_grid_args)...; underlying_grid_kwargs...) reconstructed_grid = ImmersedBoundaryGrid(reconstructed_underlying_grid, reconstructed_ib) # Test that key properties match @@ -358,11 +364,11 @@ N = 6 # test_netcdf_grid_reconstruction(gfboundary_rectilinear_grid) # test_netcdf_grid_reconstruction(gfboundary_latlon_grid) - # test_netcdf_grid_reconstruction(gfbottom_rectilinear_grid) - # test_netcdf_grid_reconstruction(gfbottom_latlon_grid) + test_netcdf_grid_reconstruction(gfbottom_rectilinear_grid) + test_netcdf_grid_reconstruction(gfbottom_latlon_grid) - # test_netcdf_grid_reconstruction(pcbottom_rectilinear_grid) - # test_netcdf_grid_reconstruction(pcbottom_latlon_grid) + test_netcdf_grid_reconstruction(pcbottom_rectilinear_grid) + test_netcdf_grid_reconstruction(pcbottom_latlon_grid) end end end