diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 0e676d25a..b2b79fb42 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -13,8 +13,6 @@ env: steps: - label: "initialize" key: "init" - concurrency: 1 - concurrency_group: 'depot/climaocean-ci' env: TEST_GROUP: "init" command: diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 2ec13b388..999c95717 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -8,7 +8,6 @@ export SimilarityTheoryTurbulentFluxes, JRA55_prescribed_atmosphere, JRA55NetCDFBackend, - ECCOMetadata, regrid_bathymetry, retrieve_bathymetry, stretched_vertical_faces, @@ -16,7 +15,10 @@ export PowerLawStretching, LinearStretching, exponential_z_faces, JRA55_field_time_series, - ECCO_field, ECCOMetadata, + ECCO_field, + ECCOMetadata, + ECCORestoring, + LinearlyTaperedPolarMask, ocean_simulation, initialize! diff --git a/src/DataWrangling/ECCO/ECCO.jl b/src/DataWrangling/ECCO/ECCO.jl index 2c56b967f..3cecdf011 100644 --- a/src/DataWrangling/ECCO/ECCO.jl +++ b/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, LinearlyTaperedPolarMask using ClimaOcean.DataWrangling: inpaint_mask! using ClimaOcean.InitialConditions: three_dimensional_regrid!, interpolate! @@ -206,14 +206,14 @@ Keyword Arguments: function inpainted_ECCO_field(metadata::ECCOMetadata; architecture = CPU(), mask = ECCO_mask(metadata, architecture), - maxiter = Inf, + inpainting = NearestNeighborInpainting(Inf), kw...) f = ECCO_field(metadata; architecture, kw...) # Make sure all values are extended properly @info "In-painting ECCO $(metadata.name)" - inpaint_mask!(f, mask; maxiter) + inpaint_mask!(f, mask; inpainting) fill_halo_regions!(f) diff --git a/src/DataWrangling/ECCO/ECCO_mask.jl b/src/DataWrangling/ECCO/ECCO_mask.jl index 0e7736228..c34784903 100644 --- a/src/DataWrangling/ECCO/ECCO_mask.jl +++ b/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)) @@ -39,3 +40,52 @@ end @inbounds mask[i, j, k] = (Tᵢ[i, j, k] == 0) end +struct LinearlyTaperedPolarMask{N, S, Z} + northern :: N + southern :: S + z :: Z +end + +""" + LinearlyTaperedPolarMask(; northern = (70, 75), + southern = (-75, -70), + z = (-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 and is equal to zero everywhere else. +The mask has the following functional form: + +```julia +n = 1 / (northern[2] - northern[1]) * (φ - northern[1]) +s = 1 / (southern[1] - southern[2]) * (φ - southern[2]) + +within_depth = (z[1] < z < z[2]) + +mask = within_depth ? max(n, s, 0) : 0 +``` +""" +function LinearlyTaperedPolarMask(; northern = (70, 75), + southern = (-75, -70), + z = (-20, 0)) + + northern[1] > northern[2] && throw(ArgumentError("Northern latitude range is invalid, northern[1] > northern[2].")) + southern[1] > southern[2] && throw(ArgumentError("Southern latitude range is invalid, southern[1] > southern[2].")) + z[1] > z[2] && throw(ArgumentError("Depth range is invalid, z[1] > z[2].")) + + return LinearlyTaperedPolarMask(northern, southern, z) +end + +@inline function (mask::LinearlyTaperedPolarMask)(φ, z) + n = 1 / (mask.northern[2] - mask.northern[1]) * (φ - mask.northern[1]) + s = 1 / (mask.southern[1] - mask.southern[2]) * (φ - mask.southern[2]) + + within_depth = (mask.z[1] < z < mask.z[2]) + + return ifelse(within_depth, max(n, s, zero(n)), zero(n)) +end + +@inline function stateindex(mask::LinearlyTaperedPolarMask, i, j, k, grid, time, loc) + LX, LY, LZ = loc + λ, φ, z = node(i, j, k, grid, LX(), LY(), LZ()) + return mask(φ, z) +end \ No newline at end of file diff --git a/src/DataWrangling/ECCO/ECCO_metadata.jl b/src/DataWrangling/ECCO/ECCO_metadata.jl index 33ebe5ec9..8d933f1f7 100644 --- a/src/DataWrangling/ECCO/ECCO_metadata.jl +++ b/src/DataWrangling/ECCO/ECCO_metadata.jl @@ -33,6 +33,10 @@ Base.show(io::IO, metadata::ECCOMetadata) = "├── version: $(metadata.version)", '\n', "└── dir: $(metadata.dir)") +Base.summary(data::ECCOMetadata{<:Any, <:ECCO2Daily}) = "Daily ECCO2 $(data.name) dataset, from $(first(data.dates)) to $(last(data.dates))" +Base.summary(data::ECCOMetadata{<:Any, <:ECCO2Monthly}) = "Monthly ECCO2 $(data.name) dataset, from $(first(data.dates)) to $(last(data.dates))" +Base.summary(data::ECCOMetadata{<:Any, <:ECCO4Monthly}) = "Monthly ECCO4 $(data.name) dataset, from $(first(data.dates)) to $(last(data.dates))" + """ ECCOMetadata(name::Symbol; dates = DateTimeProlepticGregorian(1993, 1, 1), diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 7a8ce1564..e9d6128f6 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -10,33 +10,40 @@ using JLD2 using Dates: Second using ClimaOcean: stateindex +using ClimaOcean.DataWrangling: NearestNeighborInpainting import Oceananigans.Fields: set! +import Oceananigans.Forcings: regularize_forcing import Oceananigans.OutputReaders: new_backend, update_field_time_series! @inline instantiate(T::DataType) = T() @inline instantiate(T) = T -struct ECCONetCDFBackend{N, M} <: AbstractInMemoryBackend{Int} +struct ECCONetCDFBackend{N, I, M} <: AbstractInMemoryBackend{Int} start :: Int length :: Int + inpainting :: I metadata :: M - function ECCONetCDFBackend{N}(start::Int, length::Int, metadata) where N + function ECCONetCDFBackend{N}(start::Int, length::Int, inpainting, metadata) where N M = typeof(metadata) - return new{N, M}(start, length, metadata) + I = typeof(inpainting) + return new{N, I, M}(start, length, inpainting, metadata) end end -Adapt.adapt_structure(to, b::ECCONetCDFBackend{N}) where N = ECCONetCDFBackend{N}(b.start, b.length, nothing) +Adapt.adapt_structure(to, b::ECCONetCDFBackend{N}) where N = ECCONetCDFBackend{N}(b.start, b.length, nothing, nothing) """ - ECCONetCDFBackend(length) + ECCONetCDFBackend(length; on_native_grid = false, inpainting = NearestNeighborInpainting(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, metadata; on_native_grid = false) = ECCONetCDFBackend{on_native_grid}(1, length, metadata) +ECCONetCDFBackend(length, metadata; + on_native_grid = false, + inpainting = NearestNeighborInpainting(Inf)) = ECCONetCDFBackend{on_native_grid}(1, length, inpainting, metadata) Base.length(backend::ECCONetCDFBackend) = backend.length Base.summary(backend::ECCONetCDFBackend) = string("ECCONetCDFBackend(", backend.start, ", ", backend.length, ")") @@ -44,19 +51,20 @@ Base.summary(backend::ECCONetCDFBackend) = string("ECCONetCDFBackend(", backend. const ECCONetCDFFTS{N} = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:ECCONetCDFBackend{N}} where N new_backend(b::ECCONetCDFBackend{N}, start, length) where N = - ECCONetCDFBackend{N}(start, length, b.metadata) + ECCONetCDFBackend{N}(start, length, b.inpainting, b.metadata) on_native_grid(::ECCONetCDFBackend{N}) where N = N -function set!(fts::ECCONetCDFFTS) #, metadata::ECCOMetadata=fts.backend.metadata, name::String="") +function set!(fts::ECCONetCDFFTS) backend = fts.backend - start = backend.start + start = backend.start + inpainting = backend.inpainting len = backend.length for t in start:start+len-1 # Set each element of the time-series to the associated file metadatum = @inbounds backend.metadata[t] - set!(fts[t], metadatum) + set!(fts[t], metadatum; inpainting) end fill_halo_regions!(fts) @@ -94,6 +102,7 @@ end architecture = CPU(), time_indices_in_memory = 2, time_indexing = Cyclical(), + inpainting_iterations = prod(size(metadata)), grid = nothing) Create a field time series object for ECCO data. @@ -105,24 +114,31 @@ 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()). +- `inpainting`: The inpainting algorithm to use for ECCO interpolation. For the moment, the only option is `NearestNeighborInpainting(maxiter)`, + where an average of the valid surrounding values is used `maxiter` times. - 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(), + inpainting = NearestNeighborInpainting(prod(size(metadata))), grid = nothing) # Making sure all the required individual files are downloaded download_dataset!(metadata) + if inpainting isa Int + inpainting = NearestNeighborInpainting(inpainting) + end + # ECCO data is too chunky to allow other backends backend = ECCONetCDFBackend(time_indices_in_memory, metadata; - on_native_grid = isnothing(grid)) + on_native_grid = isnothing(grid), + inpainting) loc = location(metadata) ftmp = empty_ECCO_field(first(metadata); architecture) - shortname = short_name(metadata) ECCO_native_grid = ftmp.grid fts_grid = isnothing(grid) ? ECCO_native_grid : grid @@ -160,6 +176,11 @@ oceananigans_fieldname = Dict( @inline Base.getindex(fields, i, j, k, ::UVelocity) = @inbounds fields.u[i, j, k] @inline Base.getindex(fields, i, j, k, ::VVelocity) = @inbounds fields.v[i, j, k] +Base.summary(::Temperature) = "temperature" +Base.summary(::Salinity) = "salinity" +Base.summary(::UVelocity) = "u_velocity" +Base.summary(::VVelocity) = "v_velocity" + """ struct ECCORestoring{FTS, G, M, V, N} <: Function @@ -170,41 +191,41 @@ 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 - ECCO_fts :: FTS - ECCO_grid :: G - mask :: M - variable_name :: V - λ⁻¹ :: N +struct ECCORestoring{FTS, G, M, V, N} + field_time_series :: FTS + grid :: G + mask :: M + variable_name :: V + rate :: N end Adapt.adapt_structure(to, p::ECCORestoring) = - ECCORestoring(Adapt.adapt(to, p.ECCO_fts), - Adapt.adapt(to, p.ECCO_grid), + ECCORestoring(Adapt.adapt(to, p.field_time_series), + Adapt.adapt(to, p.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) # Figure out all the inputs: time, location, and node time = Time(clock.time) - loc = location(p.ECCO_fts) + loc = location(p.field_time_series) # Retrieve the variable to force @inbounds var = fields[i, j, k, p.variable_name] - ECCO_backend = p.ECCO_fts.backend + ECCO_backend = p.field_time_series.backend native_grid = on_native_grid(ECCO_backend) - ECCO_var = get_ECCO_variable(Val(native_grid), p.ECCO_fts, i, j, k, p.ECCO_grid, grid, time) + ECCO_var = get_ECCO_variable(Val(native_grid), p.field_time_series, i, j, k, p.grid, grid, time) # 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 @@ -228,52 +249,81 @@ 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, + grid = nothing, + inpainting_iterations = prod(size(metadata))) 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. +# Positional Arguments (order does not matter): +=============================================== +- `variable_name`: The name of the variable to restore. (default: `:temperature`). + 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! +- `inpainting_iterations`: maximum number of iterations for the inpainting algorithm. (defaults to `prod(size(metadata))`) + +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::AbstractArchitecture = 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) +# Make sure we can call ECCORestoring with architecture as the first positional argument +ECCORestoring(architecture::AbstractArchitecture, variable_name::Symbol = :temperature; kw...) = + ECCORestoring(variable_name, architecture; kw...) - ECCO_fts = ECCO_field_time_series(metadata; grid, architecture, time_indices_in_memory, time_indexing) +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, + grid = nothing, + inpainting = NearestNeighborInpainting(prod(size(metadata)))) + + ECCO_fts = ECCO_field_time_series(metadata; grid, architecture, time_indices_in_memory, time_indexing, inpainting) 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 +Base.show(io::IO, p::ECCORestoring) = + print(io, "Three-dimensional restoring to ECCO data:", '\n', + "├── restored variable: ", summary(p.variable_name), '\n', + "├── restoring dataset: ", summary(p.field_time_series.backend.metadata), '\n', + "├── restoring rate: ", p.rate, '\n', + "├── mask: ", summary(p.mask), '\n', + "└── grid: ", summary(p.grid)) + +regularize_forcing(forcing::ECCORestoring, field, field_name, model_field_names) = forcing \ No newline at end of file diff --git a/src/DataWrangling/inpaint_mask.jl b/src/DataWrangling/inpaint_mask.jl index a74279348..11184a997 100644 --- a/src/DataWrangling/inpaint_mask.jl +++ b/src/DataWrangling/inpaint_mask.jl @@ -8,9 +8,20 @@ using Oceananigans.Architectures: architecture, device, GPU using KernelAbstractions: @kernel, @index +""" + NearestNeighborInpainting{M} + +A structure representing the nearest neighbor inpainting algorithm, where a missing value is +substituted with the average of the surrounding valid values. This process is repeated a maximum +of `maxiter` times or until the field is completely inpainted. +""" +struct NearestNeighborInpainting{M} + maxiter :: M +end + # Maybe we can remove this propagate field in lieu of a diffusion, # Still we'll need to do this a couple of steps on the original grid -@kernel function _propagate_field!(field, tmp_field) +@kernel function _propagate_field!(field, ::NearestNeighborInpainting, tmp_field) i, j, k = @index(Global, NTuple) @inbounds begin @@ -45,9 +56,9 @@ end propagate_horizontally!(field, ::Nothing, tmp_field=deepcopy(field); kw...) = field -function propagating(field, mask, iter, maxiter) +function propagating(field, mask, iter, inpainting::NearestNeighborInpainting) mask_sum = sum(field; condition=interior(mask)) - return isnan(mask_sum) && iter < maxiter + return isnan(mask_sum) && iter < inpainting.maxiter end @kernel function _fill_nans!(field) @@ -56,13 +67,15 @@ end end """ - propagate_horizontally!(field, mask [, tmp_field=deepcopy(field)]; max_iter = Inf) + propagate_horizontally!(inpainting, field, mask [, tmp_field=deepcopy(field)]) Horizontally propagate the values of `field` into the `mask`. In other words, cells where `mask[i, j, k] == false` are preserved, and cells where `mask[i, j, k] == true` are painted over. + +The first argument `inpainting` is the inpainting algorithm to use in the _propagate_field! step. """ -function propagate_horizontally!(field, mask, tmp_field=deepcopy(field); maxiter = Inf) +function propagate_horizontally!(inpainting::NearestNeighborInpainting, field, mask, tmp_field=deepcopy(field)) iter = 0 grid = field.grid arch = architecture(grid) @@ -73,8 +86,8 @@ function propagate_horizontally!(field, mask, tmp_field=deepcopy(field); maxiter # Need temporary field to avoid a race condition parent(tmp_field) .= parent(field) - while propagating(field, mask, iter, maxiter) - launch!(arch, grid, :xyz, _propagate_field!, field, tmp_field) + while propagating(field, mask, iter, inpainting) + launch!(arch, grid, :xyz, _propagate_field!, field, inpainting, tmp_field) launch!(arch, grid, :xyz, _substitute_values!, field, tmp_field) iter += 1 @debug "Propagate pass $iter with sum $(sum(parent(field)))" @@ -123,12 +136,17 @@ Arguments - `field`: `Field` to be inpainted. - `mask`: Boolean-valued `Field`, values where `mask[i, j, k] == true` are inpainted. - - `max_iter`: Maximum iterations for inpainting. Non-Inf values mean that - NaN's can occur within the mask. + - `inpainting`: The inpainting algorithm to use. For the moment, the only option is `NearestNeighborInpainting(maxiter)`, + where an average of the valid surrounding values is used `maxiter` times. """ -function inpaint_mask!(field, mask; maxiter = 10) +function inpaint_mask!(field, mask; inpainting = NearestNeighborInpainting(10)) + + if inpainting isa Int + inpainting = NearestNeighborInpainting(inpainting) + end + continue_downwards!(field, mask) - propagate_horizontally!(field, mask; maxiter) + propagate_horizontally!(inpainting, field, mask) return field end diff --git a/test/runtests.jl b/test/runtests.jl index 80cfaff5a..51118ad3d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,8 @@ # Common test setup file to make stand-alone tests easy include("runtests_setup.jl") +using CUDA + test_group = get(ENV, "TEST_GROUP", :all) test_group = Symbol(test_group) @@ -18,6 +20,7 @@ end if test_group == :init || test_group == :all using CUDA + CUDA.set_runtime_version!(v"12.6"; local_toolkit = true) CUDA.precompile_runtime() # Download bathymetry data diff --git a/test/test_ecco.jl b/test/test_ecco.jl index 4c1b73748..a2779d3f4 100644 --- a/test/test_ecco.jl +++ b/test/test_ecco.jl @@ -5,22 +5,31 @@ using Dates using ClimaOcean using ClimaOcean.ECCO -using ClimaOcean.ECCO: ECCO_field, metadata_path +using ClimaOcean.ECCO: ECCO_field, metadata_path, ECCO_times +using ClimaOcean.DataWrangling: NearestNeighborInpainting using Oceananigans.Grids: topology +using CFTime +using Dates + +using CUDA: @allowscalar + +start_date = DateTimeProlepticGregorian(1993, 1, 1) +end_date = DateTimeProlepticGregorian(1993, 2, 1) +dates = start_date : Month(1) : end_date + +# Inpaint only the first two cells inside the missing mask +inpainting = NearestNeighborInpainting(2) + @testset "ECCO fields utilities" begin for arch in test_architectures A = typeof(arch) @info "Testing ECCO_field on $A..." - start_date = DateTimeProlepticGregorian(1993, 1, 1) - end_date = DateTimeProlepticGregorian(1993, 4, 1) - dates = start_date : Month(1) : end_date - - temperature = ECCOMetadata(:temperature, dates) - t_restoring = ECCO_restoring_forcing(temperature; timescale = 1000.0) + temperature = ECCOMetadata(:temperature, dates, ECCO4Monthly()) + t_restoring = ECCORestoring(temperature; rate = 1 / 1000.0, inpainting) - ECCO_fts = t_restoring.func.ECCO_fts + ECCO_fts = t_restoring.field_time_series for metadatum in temperature @test isfile(metadata_path(metadatum)) @@ -37,6 +46,77 @@ using Oceananigans.Grids: topology @test Ny == size(temperature)[2] @test Nz == size(temperature)[3] @test Nt == size(temperature)[4] + + @test ECCO_fts.times[1] == ECCO_times(temperature)[1] + @test ECCO_fts.times[end] == ECCO_times(temperature)[end] + end +end + +@testset "Inpainting algorithm" begin + for arch in test_architectures + temperature = ECCOMetadata(:temperature, dates[1], ECCO4Monthly()) + + grid = LatitudeLongitudeGrid(arch, + size = (100, 100, 10), + latitude = (-75, 75), + longitude = (0, 360), + z = (-200, 0), + halo = (6, 6, 6)) + + fully_inpainted_field = CenterField(grid) + partially_inpainted_field = CenterField(grid) + + set!(fully_inpainted_field, temperature; inpainting = NearestNeighborInpainting(Inf)) + set!(partially_inpainted_field, temperature; inpainting = NearestNeighborInpainting(1)) + + fully_inpainted_interior = on_architecture(CPU(), interior(fully_inpainted_field)) + partially_inpainted_interior = on_architecture(CPU(), interior(partially_inpainted_field)) + + @test all(fully_inpainted_interior .!= 0) + @test any(partially_inpainted_interior .== 0) + end +end + +@testset "LinearlyTaperedPolarMask" begin + for arch in test_architectures + grid = LatitudeLongitudeGrid(arch; + size = (100, 100, 10), + latitude = (-75, 75), + longitude = (0, 360), + z = (-200, 0), + halo = (6, 6, 6)) + + φ₁ = @allowscalar grid.φᵃᶜᵃ[1] + φ₂ = @allowscalar grid.φᵃᶜᵃ[21] + φ₃ = @allowscalar grid.φᵃᶜᵃ[80] + φ₄ = @allowscalar grid.φᵃᶜᵃ[100] + z₁ = @allowscalar grid.zᵃᵃᶜ[6] + + mask = LinearlyTaperedPolarMask(northern = (φ₃, φ₄), + southern = (φ₁, φ₂), + z = (z₁, 0)) + + t_restoring = ECCORestoring(:temperature, arch; + dates, + mask, + rate = 1 / 1000.0, + inpainting) + + fill!(t_restoring.field_time_series[1], 1.0) + fill!(t_restoring.field_time_series[2], 1.0) + + T = CenterField(grid) + fields = (; T) + clock = Clock(; time = 0) + + @test @allowscalar t_restoring(1, 1, 10, grid, clock, fields) == t_restoring.rate + @test @allowscalar t_restoring(1, 11, 10, grid, clock, fields) == t_restoring.rate / 2 + @test @allowscalar t_restoring(1, 21, 10, grid, clock, fields) == 0 + @test @allowscalar t_restoring(1, 80, 10, grid, clock, fields) == 0 + @test @allowscalar t_restoring(1, 90, 10, grid, clock, fields) == t_restoring.rate / 2 + @test @allowscalar t_restoring(1, 100, 10, grid, clock, fields) == t_restoring.rate + @test @allowscalar t_restoring(1, 1, 5, grid, clock, fields) == 0 + @test @allowscalar t_restoring(1, 10, 5, grid, clock, fields) == 0 end end @@ -49,6 +129,34 @@ end end end +@testset "Timestepping with ECCORestoring" begin + for arch in test_architectures + + grid = LatitudeLongitudeGrid(arch; size = (10, 10, 10), + latitude = (-60, -40), + longitude = (10, 15), + z = (-200, 0), + halo = (6, 6, 6)) + + field = CenterField(grid) + + @test begin + set!(field, ECCOMetadata(:temperature)) + set!(field, ECCOMetadata(:salinity)) + true + end + + FT = ECCORestoring(:temperature, arch; dates, rate = 1 / 1000.0, inpainting) + ocean = ocean_simulation(grid; forcing = (; T = FT)) + + @test begin + time_step!(ocean) + time_step!(ocean) + true + end + end +end + @testset "Setting temperature and salinity to ECCO" begin for arch in test_architectures grid = LatitudeLongitudeGrid(size=(10, 10, 10), latitude=(-60, -40), longitude=(10, 15), z=(-200, 0), halo = (7, 7, 7))