Caching inpainting data + some updates to the one degree simulation

experiments/one_degree_simulation/one_degree_simulation.jl

@@ -1,84 +1,121 @@
 using ClimaOcean
-using ClimaOcean.ECCO: ECCO4Monthly
+using ClimaOcean.ECCO: ECCO4Monthly, NearestNeighborInpainting
 using OrthogonalSphericalShellGrids
 using Oceananigans
 using Oceananigans.Units
 using CFTime
 using Dates
 using Printf
+using CUDA: @allowscalar, device!
 
-arch = CPU()
-z = exponential_z_faces(Nz=30, depth=6000)
-Nx = 120
-Ny = 60
-Nz = length(z) - 1
+using Oceananigans.Grids: znode
 
-grid = TripolarGrid(arch; z, size = (Nx, Ny, Nz), north_poles_latitude=55, first_pole_longitude=70)
+device!(3)
+arch = GPU()
 
-@show grid
+#####
+##### Grid and Bathymetry
+#####
 
-bottom_height = regrid_bathymetry(grid;
+Nx = 360
+Ny = 180
+Nz = 100
+
+z_faces = exponential_z_faces(; Nz, depth=5000, h=34)
+
+underlying_grid = TripolarGrid(arch;
+                               size = (Nx, Ny, Nz),
+                               z = z_faces,
+                               first_pole_longitude = 70,
+                               north_poles_latitude = 55)
+
+bottom_height = regrid_bathymetry(underlying_grid;
                                   minimum_depth = 10,
-                                  interpolation_passes = 5,
-                                  major_basins = 3)
+                                  interpolation_passes = 75,
+                                  major_basins = 2)
+
+# Open Gibraltar strait 
+# TODO: find a better way to do this
+tampered_bottom_height = deepcopy(bottom_height)
+view(tampered_bottom_height, 102:103, 124, 1) .= -400
+
+grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(tampered_bottom_height))
+
+#####
+##### Closures
+#####
 
-# Closure
 gm = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=1000, κ_symmetric=1000)
-catke = Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivity()
-closure = (gm, catke)
-
-# Polar restoring
-@inline function restoring_mask(λ, φ, z, t=0)
-    ϵN = (φ - 75) / 5
-    ϵN = clamp(ϵN, zero(ϵN), one(ϵN))
-    ϵS = - (φ + 75) / 5
-    ϵS = clamp(ϵS, zero(ϵS), one(ϵS))
-    return ϵN + ϵS
-end
+catke = ClimaOcean.OceanSimulations.default_ocean_closure()
+viscous_closure = Oceananigans.TurbulenceClosures.HorizontalScalarDiffusivity(ν=10000)
+
+closure = (gm, catke, viscous_closure)
 
-restoring_rate = 1 / 1days
+#####
+##### Restoring
+#####
 
-restoring_mask_field = CenterField(grid)
-set!(restoring_mask_field, restoring_mask)
+restoring_rate  = 1 / 2days
+z_below_surface = @allowscalar znode(1, 1, grid.Nz, grid, Center(), Center(), Face())
 
-@inline sponge_layer(λ, φ, z, t, c, ω) = - restoring_mask(λ, φ, z, t) * ω * c
-Fu = Forcing(sponge_layer, field_dependencies=:u, parameters=restoring_rate)
-Fv = Forcing(sponge_layer, field_dependencies=:v, parameters=restoring_rate)
+mask = LinearlyTaperedPolarMask(southern=(-80, -70), northern=(70, 90), z=(z_below_surface, 0))
 
-dates = DateTimeProlepticGregorian(1993, 1, 1) : Month(1) : DateTimeProlepticGregorian(1994, 1, 1)
-temperature = ECCOMetadata(:temperature, dates, ECCO4Monthly())
-salinity = ECCOMetadata(:salinity, dates, ECCO4Monthly())
+dates = DateTimeProlepticGregorian(1993, 1, 1) : Month(1) : DateTimeProlepticGregorian(1993, 11, 1)
+temperature = ECCOMetadata(:temperature; dates, version=ECCO4Monthly(), dir="./")
+salinity    = ECCOMetadata(:salinity;    dates, version=ECCO4Monthly(), dir="./")
 
-FT = ECCORestoring(arch, temperature; grid, mask=restoring_mask_field, rate=restoring_rate)
-FS = ECCORestoring(arch, salinity;    grid, mask=restoring_mask_field, rate=restoring_rate)
-forcing = (T=FT, S=FS, u=Fu, v=Fv)
+# inpainting = NearestNeighborInpainting(30) should be enough to fill the gaps near bathymetry
+FT = ECCORestoring(arch, temperature; grid, mask, rate=restoring_rate, inpainting=NearestNeighborInpainting(30))
+FS = ECCORestoring(arch, salinity;    grid, mask, rate=restoring_rate, inpainting=NearestNeighborInpainting(30))
+forcing = (T=FT, S=FS)
+
+#####
+##### Ocean simulation
+##### 
 
 momentum_advection = VectorInvariant()
-tracer_advection = Centered(order=2)
+tracer_advection   = Centered(order=2)
+
+# Should we add a side drag since this is at a coarser resolution?
 ocean = ocean_simulation(grid; momentum_advection, tracer_advection,
                          closure, forcing,
                          tracers = (:T, :S, :e))
 
-set!(ocean.model,
-     T = ECCOMetadata(:temperature; dates=first(dates)),
-     S = ECCOMetadata(:salinity; dates=first(dates)))
+set!(ocean.model, T=ECCOMetadata(:temperature; dates=first(dates)),
+                  S=ECCOMetadata(:salinity;    dates=first(dates)))
+
+#####
+##### Atmospheric forcing
+#####
+
+radiation  = Radiation(arch)
+atmosphere = JRA55_prescribed_atmosphere(arch; backend=JRA55NetCDFBackend(20))
+
+#####
+##### Coupled simulation
+#####
 
-radiation = Radiation(arch)
-atmosphere = JRA55_prescribed_atmosphere(arch; backend=JRA55NetCDFBackend(41))
 sea_ice = ClimaOcean.OceanSeaIceModels.MinimumTemperatureSeaIce()
 coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
 
-simulation = Simulation(coupled_model; Δt=1, stop_iteration=10)
+simulation = Simulation(coupled_model; Δt=15minutes, stop_time=2*365days)
+
+#####
+##### Run it!
+##### 
 
 wall_time = Ref(time_ns())
 
 function progress(sim)
     ocean = sim.model.ocean
     u, v, w = ocean.model.velocities
     T = ocean.model.tracers.T
-    Tmax = maximum(T)
-    Tmin = minimum(T)
-    umax = maximum(abs, u), maximum(abs, v), maximum(abs, w)
+    Tmax = maximum(interior(T))
+    Tmin = minimum(interior(T))
+    umax = (maximum(abs, interior(u)),
+            maximum(abs, interior(v)),
+            maximum(abs, interior(w)))
+
     step_time = 1e-9 * (time_ns() - wall_time[])
 
     @info @sprintf("Time: %s, n: %d, Δt: %s, max|u|: (%.2e, %.2e, %.2e) m s⁻¹, extrema(T): (%.2f, %.2f) ᵒC, wall time: %s \n",
@@ -90,6 +127,6 @@ function progress(sim)
      return nothing
 end
 
-add_callback!(simulation, progress, IterationInterval(1))
+add_callback!(simulation, progress, IterationInterval(10))
 
 run!(simulation)

src/DataWrangling/ECCO/ECCO.jl

@@ -18,12 +18,14 @@ using Oceananigans.Utils
 
 using KernelAbstractions: @kernel, @index
 using NCDatasets
+using JLD2
 using Downloads: download
 using Dates
 using Adapt
 using Scratch
 
 download_ECCO_cache::String = ""
+
 function __init__()
     global download_ECCO_cache = @get_scratch!("ECCO")
 end
@@ -123,43 +125,68 @@ end
 """
     ECCO_field(metadata::ECCOMetadata;
                architecture = CPU(),
-               horizontal_halo = (3, 3))
-
-Retrieve the ecco field corresponding to `metadata`. 
-The data is loaded from `filename` on `architecture` with `horizontal_halo`
-in the x and y direction. The halo in the z-direction is one.
+               inpainting = nothing,
+               mask = nothing,
+               horizontal_halo = (7, 7),
+               cache_inpainted_data = false)
+
+Return a `Field` on `architecture` described by `ECCOMetadata` with
+`horizontal_halo` size.
+If not `nothing`, the `inpainting` method is used to fill the cells
+within the specified `mask`. `mask` is set to `ECCO_mask` for non-nothing
+`inpainting`.
 """
 function ECCO_field(metadata::ECCOMetadata;
                     architecture = CPU(),
-                    horizontal_halo = (7, 7))
+                    inpainting = nothing,
+                    mask = nothing,
+                    horizontal_halo = (7, 7),
+                    cache_inpainted_data = false)
+
+    # Respect user-supplied mask, but otherwise build default ECCO mask.
+    if !isnothing(inpainting) && isnothing(mask)
+        mask = ECCO_mask(metadata, architecture)
+    end
+
+    field = empty_ECCO_field(metadata; architecture, horizontal_halo)
+    inpainted_path = inpainted_metadata_path(metadata)
+
+    if !isnothing(inpainting) && isfile(inpainted_path)
+        file = jldopen(inpainted_path, "r")
+        maxiter = file["inpainting_maxiter"]
+
+        # read data if generated with the same inpainting
+        if maxiter == inpainting.maxiter
+            data = file["data"]
+            close(file)
+            copyto!(parent(field), data)
+            return field
+        end
+    end
 
     download_dataset(metadata)
     path = metadata_path(metadata)
     ds = Dataset(path)
-
     shortname = short_name(metadata)
 
     if variable_is_three_dimensional(metadata)
         data = ds[shortname][:, :, :, 1]
-
-        # The surface layer in three-dimensional ECCO fields is at `k = 1`
         data = reverse(data, dims=3)
     else
         data = ds[shortname][:, :, 1]
     end        
 
     close(ds)
-
-    field = empty_ECCO_field(metadata; architecture, horizontal_halo)
 
+    # Convert data from Union(FT, missing} to FT
     FT = eltype(field)
     data[ismissing.(data)] .= 1e10 # Artificially large number!
-    data = if location(field)[2] == Face
+    data = if location(field)[2] == Face # ?
         new_data = zeros(FT, size(field))
         new_data[:, 1:end-1, :] .= data
         new_data    
     else
-        convert.(FT, data)
+        data = Array{FT}(data)
     end
 
     # ECCO4 data is on a -180, 180 longitude grid as opposed to ECCO2 data that
@@ -173,71 +200,58 @@ function ECCO_field(metadata::ECCOMetadata;
     set!(field, data)
     fill_halo_regions!(field)
 
+    if !isnothing(inpainting)
+        # Make sure all values are extended properly
+        name = string(metadata.name)
+        date = string(metadata.dates)
+        version = summary(metadata.version)
+        @info string("Inpainting ", version, " ", name, " data from ", date, "...")
+        start_time = time_ns()
+
+        inpaint_mask!(field, mask; inpainting)
+        fill_halo_regions!(field)
+
+        elapsed = 1e-9 * (time_ns() - start_time)
+        @info string(" ... (", prettytime(elapsed), ")")
+
+        # We cache the inpainted data to avoid recomputing it
+        @root if cache_inpainted_data
+            file = jldopen(inpainted_path, "w+")
+            file["data"] = on_architecture(CPU(), parent(field))
+            file["inpainting_maxiter"] = inpainting.maxiter
+            close(file)
+        end
+    end
+
     return field
 end
 
 # Fallback
 ECCO_field(var_name::Symbol; kw...) = ECCO_field(ECCOMetadata(var_name); kw...)
 
-"""
-    inpainted_ECCO_field(metadata::ECCOMetadata;
-                         architecture = CPU(),
-                         mask = ECCO_mask(metadata, architecture),
-                         inpainting = NearestNeighborInpainting(Inf),
-                         kw...)
-
-Retrieve the ECCO field corresponding to `metadata` inpainted to fill all the missing
-values in the original dataset.
-
-Arguments
-=========
-
-- `metadata`: the metadata corresponding to the dataset.
-
-Keyword Arguments
-=================
-
-- `architecture`: either `CPU()` or `GPU()`.
-- `mask`: the mask used to inpaint the field, see [`inpaint_mask!`](@ref).
-- `inpainting`: the inpainting algorithm, see [`inpaint_mask!`](@ref). Default: `NearestNeighborInpainting(Inf)`.
-"""
-function inpainted_ECCO_field(metadata::ECCOMetadata; 
-                              architecture = CPU(),
-                              mask = ECCO_mask(metadata, architecture),
-                              inpainting = NearestNeighborInpainting(Inf),
-                              kw...)
-
-    # Make sure all values are extended properly
-    name = string(metadata.name)
-    date = string(metadata.dates)
-    version = summary(metadata.version)
-    @info string("Inpainting ", version, " ", name, " data from ", date, "...")
-    start_time = time_ns()
-
-    f = ECCO_field(metadata; architecture, kw...)
-    inpaint_mask!(f, mask; inpainting)
-    fill_halo_regions!(f)
-    elapsed = 1e-9 * (time_ns() - start_time)
-    @info string(" ... (", prettytime(elapsed), ")")
-    return f
+function inpainted_metadata_filename(metadata::ECCOMetadata)
+    original_filename = metadata_filename(metadata)
+    without_extension = original_filename[1:end-3]
+    return without_extension * "_inpainted.jld2"
 end
 
-inpainted_ECCO_field(variable_name::Symbol; kw...) = inpainted_ECCO_field(ECCOMetadata(variable_name); kw...)
-    
-function set!(field::Field, ecco_metadata::ECCOMetadata; kw...)
+inpainted_metadata_path(metadata::ECCOMetadata) = joinpath(metadata.dir, inpainted_metadata_filename(metadata))
+
+function set!(field::Field, ECCO_metadata::ECCOMetadata; kw...)
 
     # Fields initialized from ECCO
     grid = field.grid
     arch = child_architecture(grid)
-    mask = ECCO_mask(ecco_metadata, arch)
+    mask = ECCO_mask(ECCO_metadata, arch)
 
-    f = inpainted_ECCO_field(ecco_metadata; mask,
-                             architecture = arch,
-                             kw...)
+    f = ECCO_field(ECCO_metadata; mask,
+                   architecture = arch,
+                   kw...)
 
     interpolate!(field, f)
 
     return field
 end
 
 end # Module 
+

src/DataWrangling/ECCO/ECCO_metadata.jl

@@ -23,7 +23,7 @@ Metadata information for an ECCO dataset:
 - `dir`: The directory where the dataset is stored.
 """
 struct ECCOMetadata{D, V}
-    name  :: Symbol
+    name :: Symbol
     dates :: D
     version :: V
     dir :: String