Add Bathymetry metadata to allow using diverse bathymetry datasets (#518)

* Cleaning up directory tree for metadata of datasets and introducing bathymetry metadata

* New files

* Some more changes to get the metadata toi work

* Moved everything into DataWrangling

* Moved around the necessary files

* Making the ETOPO Metadata align with other datasets; and using same download_dataset() function

* Moved Bathymetry.jl out of DataWrangling

* Adding the ETOPO Metadata

* Trying to create a native grid based on ETOPO data

* Clean up

* Added the missing imports to ETOPO??

* Fixed native_grid issue

* Tried to get size(metadata) to work...

* Apply suggestions from code review

* Finalised ETOPO implementation

* drop local dir references

* bit of cleanup

* don't import if not extend

* bump minor release

* undo bump minor; bump patch release instead

* add default regrid_bathymetry w ETOPO

* add ETOPOBathymetry -> ETOPO2022

* update regrid_bathymetry(target_grid::DistributedGrid,...)

* Trying to figure out regrid_bathymetry

* Update src/DataWrangling/ETOPO/ETOPO.jl

* regrid_bathymetry downloads if needed + few tweaks

* Update src/Bathymetry.jl

* fix docstrings

* cleanup + add @root in download_dataset

* Metadatum only has one parameter

* clean imports

* better comment

* import relative

* cleanup; we need <:Any, <:Any... why?

* fix import

* add ETOPO2022 in test_downloading

* fix test

* delete

* test both methods

* add 'hidden' _regrid_bathymetry that does not download

* fix tests

* easier way to trigger download of ETOPO bathymetry

* only import what's needed

* Apply suggestions from code review

Co-authored-by: Simone Silvestri <[email protected]>

* distributed might need to download

* homogenise info msg

* debug

* debug

* debug

* comment out MultiYearJRA55; data is not available

* fix distributed_utils tests to use new topography metadata

* uncomment tests

* cleanup

* cleanup

* cleaner

* can't have nested @root

* ensure the new_grid uses the halos passed on to it

* less intepolation_passes

* better comment

---------

Co-authored-by: Navid C. Constantinou <[email protected]>
Co-authored-by: Simone Silvestri <[email protected]>

Author: 3 people

Project.toml

@@ -2,7 +2,7 @@ name = "ClimaOcean"
 uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
 license = "MIT"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.6.8"
+version = "0.6.9"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -31,8 +31,8 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
 
 [weakdeps]
-PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [extensions]

src/Bathymetry.jl

@@ -1,74 +1,36 @@
 module Bathymetry
 
-export regrid_bathymetry, retrieve_bathymetry, download_bathymetry
+export regrid_bathymetry
 
+using Downloads
 using ImageMorphology
-using ..DataWrangling: download_progress
-using Oceananigans.DistributedComputations
-
+using KernelAbstractions: @kernel, @index
 using Oceananigans
 using Oceananigans.Architectures: architecture, on_architecture
-using Oceananigans.DistributedComputations: DistributedGrid, reconstruct_global_grid, barrier!, all_reduce
-using Oceananigans.Grids: halo_size, λnodes, φnodes, x_domain, y_domain, topology
-using Oceananigans.Utils: pretty_filesize, launch!
-using Oceananigans.Fields: interpolate!
 using Oceananigans.BoundaryConditions
-using KernelAbstractions: @kernel, @index
-using JLD2
-
+using Oceananigans.DistributedComputations
+using Oceananigans.DistributedComputations: DistributedGrid, reconstruct_global_grid, all_reduce
+using Oceananigans.Fields: interpolate!
+using Oceananigans.Grids: x_domain, y_domain, topology
+using Oceananigans.Utils: launch!
 using OffsetArrays
-using ClimaOcean
-
 using NCDatasets
-using Downloads
 using Printf
 using Scratch
 
-download_bathymetry_cache::String = ""
-function __init__()
-    global download_bathymetry_cache = @get_scratch!("Bathymetry")
-end
-
-# etopo_url = "https://www.ngdc.noaa.gov/thredds/fileServer/global/ETOPO2022/60s/60s_surface_elev_netcdf/" *
-#              "ETOPO_2022_v1_60s_N90W180_surface.nc"
-
-etopo_url = "https://www.dropbox.com/scl/fi/6pwalcuuzgtpanysn4h6f/" *
-            "ETOPO_2022_v1_60s_N90W180_surface.nc?rlkey=2t7890ruyk4nd5t5eov5768lt&st=yfxsy1lu&dl=0"
-
-"""
-    download_bathymetry(; dir = download_bathymetry_cache,
-                          url = etopo_url,
-                          filename = "ETOPO_2022_v1_60s_N90W180_surface.nc")
-
-Download the bathymetry from `url` and saves it under `filename` in the directory `dir` and
-return the full filepath where the bathymetry is saved.
-"""
-function download_bathymetry(; url = etopo_url,
-                             dir = download_bathymetry_cache,
-                             filename = "ETOPO_2022_v1_60s_N90W180_surface.nc")
-
-    filepath = joinpath(dir, filename)
-
-    #TODO: embed this into a @root macro; see failed attempts in https://github.com/CliMA/ClimaOcean.jl/pull/391
-    if !isfile(filepath)
-        Downloads.download(url, filepath; progress=download_progress)
-    end
+using ..DataWrangling: Metadatum, native_grid, metadata_path, download_dataset
+using ..DataWrangling.ETOPO: ETOPO2022
 
-    return filepath
-end
+# methods specific to bathymetric datasets live within dataset modules
 
 """
-    regrid_bathymetry(target_grid;
+    regrid_bathymetry(target_grid, metadata;
                       height_above_water = nothing,
                       minimum_depth = 0,
-                      dir = download_bathymetry_cache,
-                      url = ClimaOcean.Bathymetry.etopo_url,
-                      filename = "ETOPO_2022_v1_60s_N90W180_surface.nc",
-                      interpolation_passes = 1,
-                      major_basins = 1)
+                      major_basins = 1
+                      interpolation_passes = 1)
 
-Return bathymetry associated with the NetCDF file at `path = joinpath(dir, filename)` regridded onto `target_grid`.
-If `path` does not exist, then a download is attempted from `joinpath(url, filename)`.
+Return bathymetry that corresponds to  `metadata` onto `target_grid`.
 
 Arguments
 =========
@@ -84,13 +46,6 @@ Keyword Arguments
 - `minimum_depth`: minimum depth for the shallow regions, defined as a positive value.
                    `h > - minimum_depth` is considered land. Default: 0.
 
-- `dir`: directory of the bathymetry-containing file. Default: `download_bathymetry_cache`.
-
-- `filename`: file containing bathymetric data. Must be netCDF with fields:
-  1. `lat` vector of latitude nodes
-  2. `lon` vector of longitude nodes
-  3. `z` matrix of depth values
-
 - `interpolation_passes`: regridding/interpolation passes. The bathymetry is interpolated in
                           `interpolation_passes - 1` intermediate steps. The more the interpolation
                           steps the smoother the final bathymetry becomes.
@@ -115,15 +70,27 @@ Keyword Arguments
                   the smallest basins are removed first. `major_basins = 1` retains only the largest basin.
                   If `Inf` then no basins are removed. Default: 1.
 """
-function regrid_bathymetry(target_grid;
+function regrid_bathymetry(target_grid, metadata;
                            height_above_water = nothing,
                            minimum_depth = 0,
-                           dir = download_bathymetry_cache,
-                           url = etopo_url,
-                           filename = "ETOPO_2022_v1_60s_N90W180_surface.nc",
                            interpolation_passes = 1,
                            major_basins = 1) # Allow an `Inf` number of "lakes"
 
+    download_dataset(metadata)
+
+    return _regrid_bathymetry(target_grid, metadata;
+                              height_above_water,
+                              minimum_depth,
+                              interpolation_passes,
+                              major_basins)
+end
+
+# regrid the bathymetry assuming the data is already downloaded
+function _regrid_bathymetry(target_grid, metadata;
+                            height_above_water,
+                            minimum_depth,
+                            interpolation_passes,
+                            major_basins)
     if isinteger(interpolation_passes)
         interpolation_passes = convert(Int, interpolation_passes)
     end
@@ -132,27 +99,17 @@ function regrid_bathymetry(target_grid;
         return throw(ArgumentError("interpolation_passes has to be an integer ≥ 1"))
     end
 
-    filepath = download_bathymetry(; url, dir, filename)
-    dataset = Dataset(filepath, "r")
+    arch = architecture(target_grid)
 
+    bathymetry_native_grid = native_grid(metadata, arch; halo = (10, 10, 1))
     FT = eltype(target_grid)
 
-    φ_data = dataset["lat"][:]
-    λ_data = dataset["lon"][:]
-    z_data = convert(Array{FT}, dataset["z"][:, :])
-
-    # Convert longitude from (-180, 180) to (0, 360)
-    λ_data .+= 180
-    Nhx    = size(z_data, 1)
-    z_data = circshift(z_data, (Nhx ÷ 2, 0))
+    filepath = metadata_path(metadata)
+    dataset = Dataset(filepath, "r")
 
+    z_data = convert(Array{FT}, dataset["z"][:, :])
     close(dataset)
 
-    # Diagnose target grid information
-    arch   = architecture(target_grid)
-    λ_data = λ_data |> Array{BigFloat}
-    φ_data = φ_data |> Array{BigFloat}
-
     if !isnothing(height_above_water)
         # Overwrite the height of cells above water.
         # This has an impact on reconstruction. Greater height_above_water reduces total
@@ -161,27 +118,7 @@ function regrid_bathymetry(target_grid;
         z_data[land] .= height_above_water
     end
 
-    # Infer the "native grid" of the bathymetry data and make a bathymetry field.
-    Δλ = λ_data[2] - λ_data[1]
-    Δφ = φ_data[2] - φ_data[1]
-
-    λ₁_data = convert(Float64, λ_data[1]   - Δλ / 2)
-    λ₂_data = convert(Float64, λ_data[end] + Δλ / 2)
-    φ₁_data = convert(Float64, φ_data[1]   - Δφ / 2)
-    φ₂_data = convert(Float64, φ_data[end] + Δφ / 2)
-
-    Nxn = length(λ_data)
-    Nyn = length(φ_data)
-    Nzn = 1
-
-    native_grid = LatitudeLongitudeGrid(arch, Float32;
-                                        size = (Nxn, Nyn, Nzn),
-                                        latitude  = (φ₁_data, φ₂_data),
-                                        longitude = (λ₁_data, λ₂_data),
-                                        z = (0, 1),
-                                        halo = (10, 10, 1))
-
-    native_z = Field{Center, Center, Nothing}(native_grid)
+    native_z = Field{Center, Center, Nothing}(bathymetry_native_grid)
     set!(native_z, z_data)
     fill_halo_regions!(native_z)
 
@@ -201,6 +138,57 @@ function regrid_bathymetry(target_grid;
     return target_z
 end
 
+"""
+    regrid_bathymetry(target_grid; dataset=ETOPO2022(), kw...)
+
+Regrid bathymetry from `dataset` onto `target_grid`. Default: `dataset = ETOPO2022()`.
+"""
+function regrid_bathymetry(target_grid; dataset = ETOPO2022(), kw...)
+    metadatum = Metadatum(:bottom_height; dataset)
+    return regrid_bathymetry(target_grid, metadatum; kw...)
+end
+
+# Regridding bathymetry for distributed grids, we handle the whole process
+# on just one rank, and share the results with the other processors.
+function regrid_bathymetry(target_grid::DistributedGrid, metadata;
+                           height_above_water = nothing,
+                           minimum_depth = 0,
+                           interpolation_passes = 1,
+                           major_basins = 1)
+
+    download_dataset(metadata)
+
+    global_grid = reconstruct_global_grid(target_grid)
+    global_grid = on_architecture(CPU(), global_grid)
+    arch = architecture(target_grid)
+    Nx, Ny, _ = size(global_grid)
+
+    # If all ranks open a gigantic bathymetry and the memory is
+    # shared, we could easily have OOM errors.
+    # We perform the reconstruction only on rank 0 and share the result.
+    bottom_height = if arch.local_rank == 0
+        # use regrid method that assumes data is downloaded
+        bottom_field = _regrid_bathymetry(global_grid, metadata;
+                                          height_above_water, minimum_depth, interpolation_passes, major_basins)
+        bottom_field.data[1:Nx, 1:Ny, 1]
+    else
+        zeros(Nx, Ny)
+    end
+
+    # Synchronize
+    Oceananigans.DistributedComputations.global_barrier(arch.communicator)
+
+    # Share the result (can we share SubArrays?)
+    bottom_height = all_reduce(+, bottom_height, arch)
+
+    # Partition the result
+    local_bottom_height = Field{Center, Center, Nothing}(target_grid)
+    set!(local_bottom_height, bottom_height)
+    fill_halo_regions!(local_bottom_height)
+
+    return local_bottom_height
+end
+
 @kernel function _enforce_minimum_depth!(target_z, minimum_depth)
     i, j = @index(Global, NTuple)
     z = @inbounds target_z[i, j, 1]
@@ -240,17 +228,20 @@ function interpolate_bathymetry_in_passes(native_z, target_grid;
     old_z  = native_z
     TX, TY = topology(target_grid)
 
+    Hx, Hy, Hz = Oceananigans.halo_size(native_z.grid)
+
     @info "Interpolation passes of bathymetry size $(size(old_z)) onto a $gridtype target grid of size $Nt:"
     for pass = 1:passes - 1
         new_size = (Nλ[pass], Nφ[pass], 1)
-        @info "    Performing pass $pass to size $new_size."
+        @info "    pass $pass to size $new_size"
 
         new_grid = LatitudeLongitudeGrid(architecture(target_grid), Float32,
                                          size = new_size,
                                          latitude = (latitude[1],  latitude[2]),
                                          longitude = (longitude[1], longitude[2]),
                                          z = (0, 1),
-                                         topology = (TX, TY, Bounded))
+                                         topology = (TX, TY, Bounded),
+                                         halo = (Hx, Hy, Hz))
 
         new_z = Field{Center, Center, Nothing}(new_grid)
 
@@ -266,39 +257,6 @@ function interpolate_bathymetry_in_passes(native_z, target_grid;
     return target_z
 end
 
-# Regridding bathymetry for distributed grids, we handle the whole process
-# on just one rank, and share the results with the other processors.
-function regrid_bathymetry(target_grid::DistributedGrid; kw...)
-    global_grid = reconstruct_global_grid(target_grid)
-    global_grid = on_architecture(CPU(), global_grid)
-    arch = architecture(target_grid)
-    Nx, Ny, _ = size(global_grid)
-
-    # If all ranks open a gigantic bathymetry and the memory is
-    # shared, we could easily have OOM errors.
-    # We perform the reconstruction only on rank 0 and share the result.
-    bottom_height = if arch.local_rank == 0
-        bottom_field = regrid_bathymetry(global_grid; kw...)
-        bottom_field.data[1:Nx, 1:Ny, 1]
-    else
-        zeros(Nx, Ny)
-    end
-
-    # Synchronize
-    Oceananigans.DistributedComputations.global_barrier(arch.communicator)
-
-    # Share the result (can we share SubArrays?)
-    bottom_height = all_reduce(+, bottom_height, arch)
-
-    # Partition the result
-    local_bottom_height = Field{Center, Center, Nothing}(target_grid)
-    set!(local_bottom_height, bottom_height)
-    fill_halo_regions!(local_bottom_height)
-
-    return local_bottom_height
-end
-
-
 """
     remove_minor_basins!(z_data, keep_major_basins)
 
@@ -402,39 +360,4 @@ function remove_minor_basins!(zb, keep_major_basins)
     return nothing
 end
 
-"""
-    retrieve_bathymetry(grid, filename; kw...)
-
-Retrieve the bathymetry data from a file or generate it using a grid and save it to a file.
-
-Arguments
-=========
-
-- `grid`: The grid used to generate the bathymetry data.
-- `filename`: The name of the file to read or save the bathymetry data.
-- `kw...`: Additional keyword arguments.
-
-Returns
-=======
-
-- `bottom_height`: The retrieved or generated bathymetry data.
-
-If the specified file exists, the function reads the bathymetry data from the file.
-Otherwise, it generates the bathymetry data using the provided grid and saves it to the file before returning it.
-"""
-function retrieve_bathymetry(grid, filename; kw...)
-
-    if isfile(filename)
-        bottom_height = jldopen(filename)["bathymetry"]
-    else
-        bottom_height = regrid_bathymetry(grid; kw...)
-        jldsave(filename, bathymetry = Array(interior(bottom_height)))
-    end
-
-    return bottom_height
-end
-
-retrieve_bathymetry(grid, ::Nothing; kw...) = regrid_bathymetry(grid; kw...)
-retrieve_bathymetry(grid; kw...)            = regrid_bathymetry(grid; kw...)
-
 end # module