New interface for ECCORestoring

.gitignore

@@ -1,3 +1,5 @@
+/Manifest.toml
+
 # Files generated by invoking Julia with --code-coverage
 *.jl.cov
 *.jl.*.cov

src/ClimaOcean.jl

@@ -8,15 +8,17 @@ export
     SimilarityTheoryTurbulentFluxes,
     JRA55_prescribed_atmosphere,
     JRA55NetCDFBackend,
-    ECCOMetadata,
     regrid_bathymetry,
     retrieve_bathymetry,
     stretched_vertical_faces,
     exponential_z_faces,
     PowerLawStretching, LinearStretching,
     exponential_z_faces,
     JRA55_field_time_series,
-    ECCO_field, ECCOMetadata,
+    ECCO_field, 
+    ECCOMetadata,
+    ECCORestoring,
+    LatitudinallyTaperedPolarMask,
     ocean_simulation,
     initialize!
 

src/DataWrangling/ECCO/ECCO.jl

@@ -2,7 +2,7 @@ module ECCO
 
 export ECCOMetadata, ECCO_field, ECCO_mask, adjusted_ECCO_tracers, initialize!
 export ECCO2Monthly, ECCO4Monthly, ECCO2Daily
-export ECCO_restoring_forcing
+export ECCORestoring, LatitudinallyTaperedPolarMask
 
 using ClimaOcean.DataWrangling: inpaint_mask!
 using ClimaOcean.InitialConditions: three_dimensional_regrid!, interpolate!

src/DataWrangling/ECCO/ECCO_mask.jl

@@ -1,4 +1,5 @@
 using Oceananigans.Architectures: AbstractArchitecture
+import ClimaOcean: stateindex
 
 """
     ECCO_mask(architecture = CPU(); minimum_value = Float32(-1e5))
@@ -40,3 +41,48 @@ end
     @inbounds mask[i, j, k] = (Tᵢ[i, j, k] == 0) 
 end
 
+struct LatitudinallyTaperedPolarMask{N, S, Z} 
+    northern_edges :: N
+    southern_edges :: S
+    z_edges :: Z
+end
+
+"""
+    LatitudinallyTaperedPolarMask(; northern_edges = (70,   75),
+                                    southern_edges = (-75, -70),
+                                    z_edges = (-20, 0))
+
+Build a mask that is linearly tapered in latitude inbetween the northern and southern edges.
+The mask is constant in depth between the z_edges and is equal to zero everywhere else.
+The mask has the following functional form:
+
+```julia
+n = 1 / (northern_edges[2] - northern_edges[1]) * (φ - northern_edges[1])
+s = 1 / (southern_edges[1] - southern_edges[2]) * (φ - southern_edges[2])
+
+within_depth = (z_edges[1] < z < z_edges[2])
+
+mask = within_depth ? max(n, s, 0) : 0
+```
+"""
+function LatitudinallyTaperedPolarMask(; northern_edges = (70,   75),
+                                         southern_edges = (-75, -70),
+                                         z_edges = (-20, 0))
+
+    return LatitudinallyTaperedPolarMask(northern_edges, southern_edges, z_edges)
+end
+
+@inline function (mask::LatitudinallyTaperedPolarMask)(φ, z)
+    n = 1 / (mask.northern_edges[2] - mask.northern_edges[1]) * (φ - mask.northern_edges[1])
+    s = 1 / (mask.southern_edges[1] - mask.southern_edges[2]) * (φ - mask.southern_edges[2])
+
+    within_depth = (mask.z_edges[1] < z < mask.z_edges[2])
+
+    return ifelse(within_depth, max(n, s, zero(n)), zero(n))
+end
+
+@inline function stateindex(mask::LatitudinallyTaperedPolarMask, i, j, k, grid, time, loc)
+    LX, LY, LZ = loc 
+    λ, φ, z = node(i, j, k, grid, LX(), LY(), LZ())
+    return mask(φ, z)
+end

src/DataWrangling/ECCO/ECCO_restoring.jl

@@ -12,6 +12,7 @@ using Dates: Second
 using ClimaOcean: stateindex
 
 import Oceananigans.Fields: set!
+import Oceananigans.Forcings: regularize_forcing
 import Oceananigans.OutputReaders: new_backend, update_field_time_series!
 
 @inline instantiate(T::DataType) = T()
@@ -20,43 +21,41 @@ import Oceananigans.OutputReaders: new_backend, update_field_time_series!
 struct ECCONetCDFBackend{N} <: AbstractInMemoryBackend{Int}
     start :: Int
     length :: Int
+    maxiter :: Int
 
-    ECCONetCDFBackend{N}(start::Int, length::Int) where N = new{N}(start, length)
+    ECCONetCDFBackend{N}(start::Int, length::Int, maxiter::Int) where N = new{N}(start, length, maxiter)
 end
 
 """
-    ECCONetCDFBackend(length)
+    ECCONetCDFBackend(length; on_native_grid = false, maxiter = Inf)
 
 Represents an ECCO FieldTimeSeries backed by ECCO native .nc files.
 Each time instance is stored in an individual file.
+the maxiter keyword argument is the maximum number of iterations for the inpainting algorithm.
 """
-ECCONetCDFBackend(length; on_native_grid = false) = ECCONetCDFBackend{on_native_grid}(1, length)
+ECCONetCDFBackend(length; on_native_grid = false, maxiter = Inf) = ECCONetCDFBackend{on_native_grid}(1, length, maxiter)
 
 Base.length(backend::ECCONetCDFBackend)  = backend.length
 Base.summary(backend::ECCONetCDFBackend) = string("ECCONetCDFBackend(", backend.start, ", ", backend.length, ")")
 
 const ECCONetCDFFTS{N} = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:ECCONetCDFBackend{N}} where N
 
-new_backend(::ECCONetCDFBackend{N}, start, length) where N = ECCONetCDFBackend{N}(start, length)
+new_backend(b::ECCONetCDFBackend{N}, start, length) where N = ECCONetCDFBackend{N}(start, length, b.maxiter)
 on_native_grid(::ECCONetCDFBackend{N}) where N = N
 
 function set!(fts::ECCONetCDFFTS, path::ECCOMetadata=fts.path, name::String=fts.name) 
 
     backend = fts.backend
-    start = backend.start
+    start   = backend.start
+    maxiter = backend.maxiter
 
     for t in start:start+length(backend)-1
 
         # find the file associated with the time index
         metadata = @inbounds path[t] 
 
-        arch = architecture(fts)
-        f = inpainted_ECCO_field(metadata; architecture = arch)
-        if on_native_grid(backend)
-            set!(fts[t], f)
-        else
-            interpolate!(fts[t], f)
-        end
+        # Set the field with the correct metadata
+        set!(fts[t], metadata; maxiter)
     end
 
     fill_halo_regions!(fts)
@@ -94,6 +93,7 @@ end
                            architecture = CPU(),
                            time_indices_in_memory = 2,
                            time_indexing = Cyclical(),
+                           maxiter = Inf,
                            grid = nothing)
 
 Create a field time series object for ECCO data.
@@ -105,17 +105,20 @@ Create a field time series object for ECCO data.
 - architecture: The architecture to use for computations (default: CPU()).
 - time_indices_in_memory: The number of time indices to keep in memory (default: 2).
 - time_indexing: The time indexing scheme to use (default: Cyclical()).
+- maxiter: The maximum number of iterations for the inpainting algorithm (default: Inf).
 - grid: if not a `nothing`, the ECCO data is directly interpolated on the `grid`,
 """
 function ECCO_field_time_series(metadata::ECCOMetadata;	
                                 architecture = CPU(),	
                                 time_indices_in_memory = 2,	
                                 time_indexing = Cyclical(),
+                                maxiter = Inf,
                                 grid = nothing)	
 
     # ECCO data is too chunky to allow other backends	
     backend = ECCONetCDFBackend(time_indices_in_memory; 
-                                on_native_grid = isnothing(grid))
+                                on_native_grid = isnothing(grid),
+                                maxiter)
 
     # Making sure all the required individual files are downloaded
     download_dataset!(metadata)
@@ -173,22 +176,22 @@ A struct representing ECCO restoring.
 - `ECCO_grid`: The native ECCO grid to interpolate from.
 - `mask`: A mask (could be a number, an array, a function or a field).
 - `variable_name`: The variable name of the variable that needs restoring.
-- `λ⁻¹`: The reciprocal of the restoring timescale.
+- `rate`: The reciprocal of the restoring timescale.
 """
-struct ECCORestoring{FTS, G, M, V, N} <: Function
+struct ECCORestoring{FTS, G, M, V, N} 
     ECCO_fts      :: FTS
     ECCO_grid     :: G
     mask          :: M
     variable_name :: V
-    λ⁻¹           :: N
+    rate          :: N
 end
 
 Adapt.adapt_structure(to, p::ECCORestoring) = 
     ECCORestoring(Adapt.adapt(to, p.ECCO_fts), 
                   Adapt.adapt(to, p.ECCO_grid),
                   Adapt.adapt(to, p.mask),
                   Adapt.adapt(to, p.variable_name),
-                  Adapt.adapt(to, p.λ⁻¹))
+                  Adapt.adapt(to, p.rate))
 
 @inline function (p::ECCORestoring)(i, j, k, grid, clock, fields)
 
@@ -206,8 +209,8 @@ Adapt.adapt_structure(to, p::ECCORestoring) =
 
     # Extracting the mask value at the current node
     mask = stateindex(p.mask, i, j, k, grid, clock.time, loc)
-
-    return p.λ⁻¹ * mask * (ECCO_var - var)
+    
+    return p.rate * mask * (ECCO_var - var)
 end
 
 # Differentiating between restoring done with an ECCO FTS
@@ -231,52 +234,65 @@ end
 @inline get_ECCO_variable(::Val{false}, ECCO_fts, i, j, k, ECCO_grid, grid, time) = @inbounds ECCO_fts[i, j, k, time]
 
 """
-    ECCO_restoring_forcing(metadata::ECCOMetadata;
-                            architecture = CPU(), 
-                            backend = ECCONetCDFBackend(2),
-                            time_indexing = Cyclical(),
-                            mask = 1,
-                            timescale = 5days)
+    ECCORestoring(variable_name::Symbol, [architecture = CPU()]; 
+                  version = ECCO4Monthly(),
+                  dates = all_ECCO_dates(version), 
+                  time_indices_in_memory = 2, # Not more than this if we want to use GPU!
+                  time_indexing = Cyclical(),
+                  mask = 1,
+                  rate = 1 / 20days,
+                  grid = nothing,
+                  maxiter = Inf)
 
 Create a restoring forcing term that restores to values stored in an ECCO field time series.
 
 # Arguments:
 =============
-- `metadata`: The metadata for the ECCO field time series.
+- `variable_name`: The name of the variable to restore. The choice is between 
+                   :temperature, :salinity, :u_velocity, :v_velocity, :sea_ice_thickness, :sea_ice_area_fraction.
+- `architecture`: The architecture. Typically `CPU` or `GPU`. Default is `CPU`.
 
 # Keyword Arguments:
 ====================
-- `architecture`: The architecture. Typically `CPU` or `GPU`
+- `version`: The version of the ECCO dataset. Default is `ECCO4Monthly()`.
+- `dates`: The dates to use for the ECCO dataset. Default is `all_ECCO_dates(version)`.
 - `time_indices_in_memory`: The number of time indices to keep in memory. trade-off between performance
                             and memory footprint.    
-- `time_indexing`: The time indexing scheme for the field time series, see [`FieldTimeSeries`](@ref)
+- `time_indexing`: The time indexing scheme for the field time series≥
 - `mask`: The mask value. Can be a function of `(x, y, z, time)`, an array or a number
-- `timescale`: The restoring timescale.
+- `rate`: The restoring rate in s⁻¹.
+- `time_indices_in_memory = 2, # Not more than this if we want to use GPU!
+- `maxiter`: maximum number of iterations for the inpainting algorithm. (defaults to `Inf`)
+
+It is possible to also pass an `ECCOMetadata` type as the first argument without the need for the 
+`variable_name` argument and the `version` and `dates` keyword arguments.
 """
-function ECCO_restoring_forcing(variable_name::Symbol, version=ECCO4Monthly(); kw...) 
-     metadata = ECCOMetadata(variable_name, all_ECCO_dates(version), version)
-    return ECCO_restoring_forcing(metadata; kw...)
+function ECCORestoring(variable_name::Symbol, architecture = CPU(); 
+                       version = ECCO4Monthly(),
+                       dates = all_ECCO_dates(version), 
+                       kw...) 
+
+     metadata = ECCOMetadata(variable_name, dates, version)
+    return ECCORestoring(metadata; architecture, kw...)
 end
 
-function ECCO_restoring_forcing(metadata::ECCOMetadata;
-                                architecture = CPU(), 
-                                time_indices_in_memory = 2, # Not more than this if we want to use GPU!
-                                time_indexing = Cyclical(),
-                                mask = 1,
-                                timescale = 20days,
-                                grid = nothing)
+function ECCORestoring(metadata::ECCOMetadata;
+                       architecture = CPU(), 
+                       time_indices_in_memory = 2, # Not more than this if we want to use GPU!
+                       time_indexing = Cyclical(),
+                       mask = 1,
+                       rate = 1 / 20days,
+                       grid = nothing.
+                       maxiter = Inf)
 
-    ECCO_fts  = ECCO_field_time_series(metadata; grid, architecture, time_indices_in_memory, time_indexing)                  
+    ECCO_fts  = ECCO_field_time_series(metadata; grid, architecture, time_indices_in_memory, time_indexing, maxiter)                  
     ECCO_grid = ECCO_fts.grid
 
     # Grab the correct Oceananigans field to restore
     variable_name = metadata.name
     field_name = oceananigans_fieldname[variable_name]
-    ECCO_restoring = ECCORestoring(ECCO_fts, ECCO_grid, mask, field_name, 1 / timescale)
-
-    # Defining the forcing that depends on the restoring field.
-    restoring_forcing = Forcing(ECCO_restoring; discrete_form = true)
 
-    return restoring_forcing
+    return ECCORestoring(ECCO_fts, ECCO_grid, mask, field_name, rate)
 end
 
+regularize_forcing(forcing::ECCORestoring, field, field_name, model_field_names) = forcing